Add duration and timestamp to well-known types

Author: AdrienVannson

src/betterproto2_compiler/known_types/__init__.py

@@ -1,10 +1,14 @@
 from collections.abc import Callable
 
 from .any import Any
+from .duration import Duration
+from .timestamp import Timestamp
 
 # For each (package, message name), lists the methods that should be added to the message definition.
 # The source code of the method is read from the `known_types` folder. If imports are needed, they can be directly added
 # 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", "Duration"): [Duration.from_timedelta, Duration.to_timedelta, Duration.delta_to_json],
 }

src/betterproto2_compiler/known_types/duration.py

@@ -0,0 +1,25 @@
+import datetime
+
+from betterproto2.lib.std.google.protobuf import Duration as VanillaDuration
+
+
+class Duration(VanillaDuration):
+    @classmethod
+    def from_timedelta(
+        cls, delta: datetime.timedelta, *, _1_microsecond: datetime.timedelta = datetime.timedelta(microseconds=1)
+    ) -> "Duration":
+        total_ms = delta // _1_microsecond
+        seconds = int(total_ms / 1e6)
+        nanos = int((total_ms % 1e6) * 1e3)
+        return cls(seconds, nanos)
+
+    def to_timedelta(self) -> datetime.timedelta:
+        return datetime.timedelta(seconds=self.seconds, microseconds=self.nanos / 1e3)
+
+    @staticmethod
+    def delta_to_json(delta: datetime.timedelta) -> str:
+        parts = str(delta.total_seconds()).split(".")
+        if len(parts) > 1:
+            while len(parts[1]) not in (3, 6, 9):
+                parts[1] = f"{parts[1]}0"
+        return f"{'.'.join(parts)}s"

src/betterproto2_compiler/known_types/timestamp.py

@@ -0,0 +1,45 @@
+import datetime
+
+from betterproto2.lib.std.google.protobuf import Timestamp as VanillaTimestamp
+
+
+class Timestamp(VanillaTimestamp):
+    @classmethod
+    def from_datetime(cls, dt: datetime.datetime) -> "Timestamp":
+        # manual epoch offset calulation to avoid rounding errors,
+        # to support negative timestamps (before 1970) and skirt
+        # around datetime bugs (apparently 0 isn't a year in [0, 9999]??)
+        offset = dt - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)
+        # below is the same as timedelta.total_seconds() but without dividing by 1e6
+        # so we end up with microseconds as integers instead of seconds as float
+        offset_us = (offset.days * 24 * 60 * 60 + offset.seconds) * 10**6 + offset.microseconds
+        seconds, us = divmod(offset_us, 10**6)
+        return cls(seconds, us * 1000)
+
+    def to_datetime(self) -> datetime.datetime:
+        # datetime.fromtimestamp() expects a timestamp in seconds, not microseconds
+        # if we pass it as a floating point number, we will run into rounding errors
+        # see also #407
+        offset = datetime.timedelta(seconds=self.seconds, microseconds=self.nanos // 1000)
+        return datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + offset
+
+    @staticmethod
+    def timestamp_to_json(dt: datetime.datetime) -> str:
+        nanos = dt.microsecond * 1e3
+        if dt.tzinfo is not None:
+            # change timezone aware datetime objects to utc
+            dt = dt.astimezone(datetime.timezone.utc)
+        copy = dt.replace(microsecond=0, tzinfo=None)
+        result = copy.isoformat()
+        if (nanos % 1e9) == 0:
+            # If there are 0 fractional digits, the fractional
+            # point '.' should be omitted when serializing.
+            return f"{result}Z"
+        if (nanos % 1e6) == 0:
+            # Serialize 3 fractional digits.
+            return f"{result}.{int(nanos // 1e6) :03d}Z"
+        if (nanos % 1e3) == 0:
+            # Serialize 6 fractional digits.
+            return f"{result}.{int(nanos // 1e3) :06d}Z"
+        # Serialize 9 fractional digits.
+        return f"{result}.{nanos:09d}"