Wrap timestamp and duration

Author: AdrienVannson

src/betterproto2_compiler/compile/importing.py

@@ -80,7 +80,7 @@ def get_type_reference(
     Return a Python type name for a proto type reference. Adds the import if
     necessary. Unwraps well known type if required.
     """
-    if unwrap:
+    if unwrap:  # TODO don't hardcode
         if source_type == ".google.protobuf.Duration":
             return "datetime.timedelta"
 

src/betterproto2_compiler/known_types/__init__.py

@@ -10,7 +10,13 @@
 # to the template file: they will automatically be removed if not necessary.
 KNOWN_METHODS: dict[tuple[str, str], list[Callable]] = {
     ("google.protobuf", "Any"): [Any.pack, Any.unpack, Any.to_dict],
-    ("google.protobuf", "Timestamp"): [Timestamp.from_datetime, Timestamp.to_datetime, Timestamp.timestamp_to_json],
+    ("google.protobuf", "Timestamp"): [
+        Timestamp.from_datetime,
+        Timestamp.to_datetime,
+        Timestamp.timestamp_to_json,
+        Timestamp.from_dict,
+        Timestamp.to_wrapped,
+    ],
     ("google.protobuf", "Duration"): [Duration.from_timedelta, Duration.to_timedelta, Duration.delta_to_json],
     ("google.protobuf", "BoolValue"): [BoolValue.from_wrapped, BoolValue.to_wrapped],
     ("google.protobuf", "StringValue"): [StringValue.from_wrapped, StringValue.to_wrapped],
@@ -20,4 +26,6 @@
 WRAPPED_TYPES: dict[tuple[str, str], str] = {
     ("google.protobuf", "BoolValue"): "bool",
     ("google.protobuf", "StringValue"): "str",
+    ("google.protobuf", "Timestamp"): "datetime.datetime",
+    ("google.protobuf", "Duration"): "datetime.timedelta",
 }

src/betterproto2_compiler/known_types/timestamp.py

@@ -1,5 +1,7 @@
 import datetime
 
+import dateutil.parser
+
 from betterproto2_compiler.lib.google.protobuf import Timestamp as VanillaTimestamp
 
 
@@ -43,3 +45,16 @@ def timestamp_to_json(dt: datetime.datetime) -> str:
             return f"{result}.{int(nanos // 1e3):06d}Z"
         # Serialize 9 fractional digits.
         return f"{result}.{nanos:09d}"
+
+    # TODO typing
+    @classmethod
+    def from_dict(cls, value):
+        if isinstance(value, str):
+            dt = dateutil.parser.isoparse(value)
+            dt = dt.astimezone(datetime.timezone.utc)
+            return Timestamp.from_datetime(dt)
+
+        return super().from_dict(value)
+
+    def to_wrapped(self) -> datetime.datetime:
+        return self.to_datetime()

src/betterproto2_compiler/plugin/models.py

@@ -317,13 +317,23 @@ def get_field_string(self) -> str:
     def betterproto_field_args(self) -> list[str]:
         args = []
 
-        if self.proto_obj.type == FieldDescriptorProtoType.TYPE_MESSAGE:
+        if self.field_type == FieldDescriptorProtoType.TYPE_MESSAGE:
             type_package, type_name = parse_source_type_name(self.proto_obj.type_name, self.output_file.parent_request)
 
             if (type_package, type_name) in WRAPPED_TYPES:
+                unwrap_type = get_type_reference(
+                    package=self.output_file.package,
+                    imports=self.output_file.imports_end,
+                    source_type=self.proto_obj.type_name,
+                    request=self.output_file.parent_request,
+                    settings=self.output_file.settings,
+                    unwrap=False,
+                )
+
                 # Without the lambda function, the type is evaluated right away, which fails since the corresponding
                 # import is placed at the end of the file to avoid circular imports.
                 args.append(f"wrap=lambda: {self.py_type}")
+                args.append(f"unwrap=lambda: {unwrap_type}")
 
         if self.optional:
             args.append("optional=True")

src/betterproto2_compiler/templates/header.py.j2

@@ -20,6 +20,7 @@ __all__ = (
 
 import builtins
 import datetime
+import dateutil.parser
 import warnings
 from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator
 import typing