Support durations

Author: AdrienVannson

src/betterproto2_compiler/known_types/__init__.py

@@ -19,7 +19,15 @@
         Timestamp.from_wrapped,
         Timestamp.to_wrapped,
     ],
-    ("google.protobuf", "Duration"): [Duration.from_timedelta, Duration.to_timedelta, Duration.delta_to_json],
+    ("google.protobuf", "Duration"): [
+        Duration.from_timedelta,
+        Duration.to_timedelta,
+        Duration.delta_to_json,
+        Duration.from_dict,
+        Duration.to_dict,
+        Duration.from_wrapped,
+        Duration.to_wrapped,
+    ],
     ("google.protobuf", "BoolValue"): [BoolValue.from_wrapped, BoolValue.to_wrapped],
     ("google.protobuf", "StringValue"): [StringValue.from_wrapped, StringValue.to_wrapped],
 }

src/betterproto2_compiler/known_types/duration.py

@@ -1,4 +1,8 @@
 import datetime
+import re
+import typing
+
+import betterproto2
 
 from betterproto2_compiler.lib.google.protobuf import Duration as VanillaDuration
 
@@ -23,3 +27,44 @@ def delta_to_json(delta: datetime.timedelta) -> str:
             while len(parts[1]) not in (3, 6, 9):
                 parts[1] = f"{parts[1]}0"
         return f"{'.'.join(parts)}s"
+
+    # TODO typing
+    @classmethod
+    def from_dict(cls, value):
+        if isinstance(value, str):
+            if not re.match(r"^\d+(\.\d+)?s$", value):
+                raise ValueError(f"Invalid duration string: {value}")
+
+            seconds = float(value[:-1])
+            return Duration(seconds=int(seconds), nanos=int((seconds - int(seconds)) * 1e9))
+
+        return super().from_dict(value)
+
+    # TODO typing
+    def to_dict(
+        self,
+        *,
+        output_format: betterproto2.OutputFormat = betterproto2.OutputFormat.PROTO_JSON,
+        casing: betterproto2.Casing = betterproto2.Casing.CAMEL,
+        include_default_values: bool = False,
+    ) -> dict[str, typing.Any] | typing.Any:
+        # If the output format is PYTHON, we should have kept the wraped type without building the real class
+        assert output_format == betterproto2.OutputFormat.PROTO_JSON
+
+        assert 0 <= self.nanos < 1e9
+
+        if self.nanos == 0:
+            return f"{self.seconds}s"
+
+        nanos = f"{self.nanos:09d}".rstrip("0")
+        if len(nanos) < 3:
+            nanos += "0" * (3 - len(nanos))
+
+        return f"{self.seconds}.{nanos}s"
+
+    @staticmethod
+    def from_wrapped(wrapped: datetime.timedelta) -> "Duration":
+        return Duration.from_timedelta(wrapped)
+
+    def to_wrapped(self) -> datetime.timedelta:
+        return self.to_timedelta()

src/betterproto2_compiler/known_types/timestamp.py

@@ -71,7 +71,6 @@ def to_dict(
         casing: betterproto2.Casing = betterproto2.Casing.CAMEL,
         include_default_values: bool = False,
     ) -> dict[str, typing.Any] | typing.Any:
-        print("ok")
         # If the output format is PYTHON, we should have kept the wraped type without building the real class
         assert output_format == betterproto2.OutputFormat.PROTO_JSON
 

src/betterproto2_compiler/templates/header.py.j2

@@ -18,6 +18,7 @@ __all__ = (
     {%- endfor -%}
 )
 
+import re
 import builtins
 import datetime
 import dateutil.parser