Implement typing for proto Timestamp/Duration assignments.

PiperOrigin-RevId: 730336339

Author: anandolee

python/google/protobuf/internal/duration_test.py

@@ -12,6 +12,7 @@
 import unittest
 
 from google.protobuf import duration
+from google.protobuf.internal import well_known_types_test_pb2
 
 from google.protobuf import duration_pb2
 
@@ -57,6 +58,13 @@ def test_duration_timedelta(self):
     converted_message = duration.from_timedelta(td)
     self.assertEqual(message, converted_message)
 
+  def test_duration_construction(self):
+    expected_td = datetime.timedelta(microseconds=123)
+    message = well_known_types_test_pb2.WKTMessage(
+        optional_duration=expected_td
+    )
+    self.assertEqual(expected_td, message.optional_duration.ToTimedelta())
+
 
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/internal/timestamp_test.py

@@ -59,6 +59,11 @@ def test_timestamp_datetime(self):
         naive_utc_epoch, timestamp.to_datetime(message.optional_timestamp)  # pytype: disable=wrong-arg-types
     )
 
+  def test_timstamp_construction(self):
+    message = well_known_types_test_pb2.WKTMessage(
+        optional_timestamp=datetime.datetime.today()
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

src/google/protobuf/compiler/python/pyi_generator.cc

@@ -396,8 +396,18 @@ std::string PyiGenerator::GetFieldType(
   return "";
 }
 
-void PyiGenerator::PrintMessage(
-    const Descriptor& message_descriptor, bool is_nested) const {
+std::string PyiGenerator::ExtraInitTypes(const Descriptor& msg_des) const {
+  if (msg_des.full_name() == "google.protobuf.Timestamp") {
+    return "datetime.datetime, ";
+  } else if (msg_des.full_name() == "google.protobuf.Duration") {
+    return "datetime.timedelta, ";
+  } else {
+    return "";
+  }
+}
+
+void PyiGenerator::PrintMessage(const Descriptor& message_descriptor,
+                                bool is_nested) const {
   if (!is_nested) {
     printer_->Print("\n");
   }
@@ -528,9 +538,11 @@ void PyiGenerator::PrintMessage(
         printer_->Print("_Iterable[");
       }
       if (field_des->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE) {
-        printer_->Print(
-            "_Union[$type_name$, _Mapping]", "type_name",
-            GetFieldType(*field_des, message_descriptor));
+        const auto& extra_init_types =
+            ExtraInitTypes(*field_des->message_type());
+        printer_->Print("_Union[$extra_init_types$$type_name$, _Mapping]",
+                        "extra_init_types", extra_init_types, "type_name",
+                        GetFieldType(*field_des, message_descriptor));
       } else {
         if (field_des->cpp_type() == FieldDescriptor::CPPTYPE_ENUM) {
           printer_->Print("_Union[$type_name$, str]", "type_name",

src/google/protobuf/compiler/python/pyi_generator.h

@@ -83,6 +83,7 @@ class PROTOC_EXPORT PyiGenerator : public google::protobuf::compiler::CodeGenera
   std::string ModuleLevelName(const DescriptorT& descriptor) const;
   std::string PublicPackage() const;
   std::string InternalPackage() const;
+  std::string ExtraInitTypes(const Descriptor& msg_des) const;
 
   bool opensource_runtime_ = true;