Preserve nanosecond precision for Timestamp

README.md

@@ -14,7 +14,7 @@ This project aims to provide an improved experience when using Protobuf / gRPC i
   - Enums
   - Dataclasses
   - `async`/`await`
-  - Timezone-aware `datetime` and `timedelta` objects
+  - Timezone-aware `datetime` and `timedelta` objects, including nanosecond-precision `Timestamp` values
   - Relative imports
   - Mypy type checking
 - [Pydantic Models](https://docs.pydantic.dev/) generation

betterproto2/docs/getting-started.md

@@ -225,6 +225,12 @@ Message objects include `betterproto.Message.to_json` and
 `betterproto.Message.to_dict`, `betterproto.Message.from_dict` for
 converting back and forth from JSON serializable dicts.
 
+`google.protobuf.Timestamp` fields use timezone-aware `datetime.datetime`
+values. When binary or JSON data contains sub-microsecond precision,
+betterproto2 preserves it by returning a `betterproto2.nano_datetime.NanoDatetime`,
+which is a `datetime.datetime` subclass. Timestamp JSON accepts and emits
+RFC 3339 strings with up to 9 fractional second digits.
+
 For compatibility the default is to convert field names to
 `betterproto.Casing.CAMEL`. You can control this behavior by passing a
 different casing value, e.g:

betterproto2/docs/index.md

@@ -12,5 +12,6 @@ Python code, using modern language features.
 
 - Generated messages are both binary & JSON serializable
 - Messages use relevant python types, e.g. ``Enum``, ``datetime`` and ``timedelta`` objects
+- ``Timestamp`` values preserve nanosecond precision when binary or JSON data contains it
 - ``async``/``await`` support for gRPC Clients and Servers
 - Generates modern, readable, idiomatic python code

betterproto2/src/betterproto2/nano_datetime.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import datetime
+import re
+from typing import Any
+
+import dateutil.parser
+from typing_extensions import Self
+
+_UTC = datetime.timezone.utc
+_TIMESTAMP_ZERO = datetime.datetime(1970, 1, 1, tzinfo=_UTC)
+_NANOS_PER_MICROSECOND = 1000
+_MICROS_PER_SECOND = 10**6
+_TIMESTAMP_RE = re.compile(
+    r"^"
+    r"(?P<date>\d{4}-\d{2}-\d{2})"
+    r"T"
+    r"(?P<time>\d{2}:\d{2}:\d{2})"
+    r"(?:\.(?P<fraction>\d{1,9}))?"
+    r"(?P<tz>Z|[+-]\d{2}:\d{2})"
+    r"$"
+)
+
+
+class NanoDatetime(datetime.datetime):
+    """A datetime that carries Timestamp sub-microsecond precision."""
+
+    __slots__ = ("_nanosecond_remainder",)
+
+    def __new__(
+        cls,
+        *args: Any,
+        nanosecond_remainder: int = 0,
+        **kwargs: Any,
+    ) -> Self:
+        if not 0 <= nanosecond_remainder < _NANOS_PER_MICROSECOND:
+            raise ValueError("nanosecond_remainder must be in range 0..999")
+
+        instance = super().__new__(cls, *args, **kwargs)
+        object.__setattr__(instance, "_nanosecond_remainder", nanosecond_remainder)
+        return instance
+
+    @property
+    def nanosecond_remainder(self) -> int:
+        return getattr(self, "_nanosecond_remainder", 0)
+
+    @property
+    def total_nanoseconds(self) -> int:
+        return self.microsecond * _NANOS_PER_MICROSECOND + self.nanosecond_remainder
+
+    def replace(
+        self, *args: Any, nanosecond_remainder: int | None = None, **kwargs: Any
+    ) -> NanoDatetime:
+        if nanosecond_remainder is None:
+            nanosecond_remainder = self.nanosecond_remainder
+
+        replaced = super().replace(*args, **kwargs)
+        return _datetime_to_nano_datetime(replaced, nanosecond_remainder)
+
+    def __repr__(self) -> str:
+        base = super().__repr__()
+        return f"{base[:-1]}, nanosecond_remainder={self.nanosecond_remainder})"
+
+    def __eq__(self, other: Any) -> bool:
+        equal = super().__eq__(other)
+        if equal is NotImplemented or not isinstance(other, datetime.datetime):
+            return equal
+
+        return bool(equal) and self.nanosecond_remainder == _datetime_nanosecond_remainder(other)
+
+    def __ne__(self, other: Any) -> bool:
+        equal = self.__eq__(other)
+        if equal is NotImplemented:
+            return equal
+
+        return not equal
+
+    def __hash__(self) -> int:
+        if self.nanosecond_remainder == 0:
+            return super().__hash__()
+
+        return hash((super().__hash__(), self.nanosecond_remainder))
+
+    @staticmethod
+    def to_timestamp(dt: datetime.datetime) -> tuple[int, int]:
+        if not dt.tzinfo:
+            raise ValueError("datetime must be timezone aware")
+
+        nanos = _datetime_total_nanoseconds(dt)
+        dt = dt.astimezone(_UTC)
+
+        offset = dt - _TIMESTAMP_ZERO
+        offset_us = (offset.days * 24 * 60 * 60 + offset.seconds) * _MICROS_PER_SECOND + offset.microseconds
+        seconds, _ = divmod(offset_us, _MICROS_PER_SECOND)
+        return seconds, nanos
+
+    @staticmethod
+    def from_timestamp(seconds: int, nanos: int) -> NanoDatetime:
+        micros, nanosecond_remainder = divmod(nanos, _NANOS_PER_MICROSECOND)
+        offset = datetime.timedelta(seconds=seconds, microseconds=micros)
+        dt = _TIMESTAMP_ZERO + offset
+        return _datetime_to_nano_datetime(dt, nanosecond_remainder)
+
+    @staticmethod
+    def to_json(dt: datetime.datetime) -> str:
+        seconds, nanos = NanoDatetime.to_timestamp(dt)
+        dt = NanoDatetime.from_timestamp(seconds, nanos)
+        result = dt.replace(microsecond=0, nanosecond_remainder=0, tzinfo=None).isoformat()
+
+        if nanos == 0:
+            return f"{result}Z"
+        if nanos % 1_000_000 == 0:
+            return f"{result}.{nanos // 1_000_000:03d}Z"
+        if nanos % 1_000 == 0:
+            return f"{result}.{nanos // 1_000:06d}Z"
+
+        return f"{result}.{nanos:09d}Z"
+
+    @staticmethod
+    def from_rfc3339(value: str) -> datetime.datetime:
+        match = _TIMESTAMP_RE.match(value)
+        if match:
+            tz = "+00:00" if match.group("tz") == "Z" else match.group("tz")
+            dt = datetime.datetime.fromisoformat(f"{match.group('date')}T{match.group('time')}{tz}")
+
+            fraction = (match.group("fraction") or "").ljust(9, "0")
+            total_nanos = int(fraction) if fraction else 0
+            micros, nanosecond_remainder = divmod(total_nanos, _NANOS_PER_MICROSECOND)
+
+            dt = dt.replace(microsecond=micros).astimezone(_UTC)
+            return _datetime_to_nano_datetime(dt, nanosecond_remainder)
+
+        dt = dateutil.parser.isoparse(value)
+        return dt.astimezone(_UTC)
+
+
+def _datetime_nanosecond_remainder(dt: datetime.datetime) -> int:
+    if isinstance(dt, NanoDatetime):
+        return dt.nanosecond_remainder
+
+    return 0
+
+
+def _datetime_total_nanoseconds(dt: datetime.datetime) -> int:
+    if isinstance(dt, NanoDatetime):
+        return dt.total_nanoseconds
+
+    return dt.microsecond * _NANOS_PER_MICROSECOND
+
+
+def _datetime_to_nano_datetime(dt: datetime.datetime, nanosecond_remainder: int = 0) -> NanoDatetime:
+    return NanoDatetime(
+        dt.year,
+        dt.month,
+        dt.day,
+        dt.hour,
+        dt.minute,
+        dt.second,
+        dt.microsecond,
+        tzinfo=dt.tzinfo,
+        fold=dt.fold,
+        nanosecond_remainder=nanosecond_remainder,
+    )

betterproto2/tests/test_nano_datetime.py

@@ -0,0 +1,83 @@
+from datetime import datetime, timezone
+
+import pytest
+
+from betterproto2.nano_datetime import NanoDatetime
+
+
+@pytest.mark.parametrize("nanosecond_remainder", [-1, 1000])
+def test_nanosecond_remainder_must_be_sub_microsecond(nanosecond_remainder: int) -> None:
+    with pytest.raises(ValueError, match="nanosecond_remainder"):
+        NanoDatetime(2024, 1, 2, tzinfo=timezone.utc, nanosecond_remainder=nanosecond_remainder)
+
+
+def test_replace_preserves_nanosecond_remainder_by_default() -> None:
+    dt = NanoDatetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc, nanosecond_remainder=789)
+
+    replaced = dt.replace(second=6)
+
+    assert replaced == NanoDatetime(
+        2024,
+        1,
+        2,
+        3,
+        4,
+        6,
+        123456,
+        tzinfo=timezone.utc,
+        nanosecond_remainder=789,
+    )
+
+
+def test_replace_can_override_nanosecond_remainder() -> None:
+    dt = NanoDatetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc, nanosecond_remainder=789)
+
+    replaced = dt.replace(nanosecond_remainder=123)
+
+    assert replaced.nanosecond_remainder == 123
+    assert replaced.total_nanoseconds == 123456123
+
+
+def test_equality_and_hash_include_nanosecond_remainder() -> None:
+    base = datetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc)
+    zero_remainder = NanoDatetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc)
+    one_remainder = NanoDatetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc, nanosecond_remainder=1)
+    two_remainder = NanoDatetime(2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc, nanosecond_remainder=2)
+
+    assert zero_remainder == base
+    assert hash(zero_remainder) == hash(base)
+    assert one_remainder != base
+    assert one_remainder != two_remainder
+    assert len({base, zero_remainder, one_remainder, two_remainder}) == 3
+
+
+@pytest.mark.parametrize(
+    ("dt", "expected"),
+    [
+        (NanoDatetime(1970, 1, 1, tzinfo=timezone.utc), "1970-01-01T00:00:00Z"),
+        (NanoDatetime(1970, 1, 1, 0, 0, 0, 123000, tzinfo=timezone.utc), "1970-01-01T00:00:00.123Z"),
+        (NanoDatetime(1970, 1, 1, 0, 0, 0, 123456, tzinfo=timezone.utc), "1970-01-01T00:00:00.123456Z"),
+        (
+            NanoDatetime(1970, 1, 1, 0, 0, 0, 123456, tzinfo=timezone.utc, nanosecond_remainder=789),
+            "1970-01-01T00:00:00.123456789Z",
+        ),
+    ],
+)
+def test_to_json_uses_required_fraction_widths(dt: NanoDatetime, expected: str) -> None:
+    assert NanoDatetime.to_json(dt) == expected
+
+
+@pytest.mark.parametrize(
+    ("value", "total_nanoseconds"),
+    [
+        ("1970-01-01T00:00:00.1Z", 100000000),
+        ("1970-01-01T00:00:00.1234Z", 123400000),
+        ("1970-01-01T00:00:00.12345678Z", 123456780),
+        ("1970-01-01T00:00:00.123456789Z", 123456789),
+    ],
+)
+def test_from_rfc3339_accepts_subsecond_fraction_widths(value: str, total_nanoseconds: int) -> None:
+    dt = NanoDatetime.from_rfc3339(value)
+
+    assert isinstance(dt, NanoDatetime)
+    assert dt.total_nanoseconds == total_nanoseconds

betterproto2/tests/test_timestamp.py

@@ -5,6 +5,7 @@
 
 import pytest
 
+from betterproto2.nano_datetime import NanoDatetime
 from tests.outputs.google.google.protobuf import Timestamp
 
 
@@ -34,3 +35,49 @@ def test_invalid_datetime():
     """
     with pytest.raises(ValueError):
         Timestamp.from_datetime(datetime.now())
+
+
+def test_timestamp_to_datetime_preserves_nanoseconds():
+    ts = Timestamp(seconds=1, nanos=123456789)
+
+    dt = ts.to_datetime()
+
+    assert isinstance(dt, NanoDatetime)
+    assert dt.microsecond == 123456
+    assert dt.nanosecond_remainder == 789
+    assert dt.total_nanoseconds == 123456789
+    assert Timestamp.from_datetime(dt) == ts
+
+
+def test_timestamp_to_datetime_preserves_negative_nanoseconds():
+    ts = Timestamp(seconds=-1, nanos=999999999)
+
+    dt = ts.to_datetime()
+
+    assert isinstance(dt, NanoDatetime)
+    assert dt == NanoDatetime(
+        1969,
+        12,
+        31,
+        23,
+        59,
+        59,
+        999999,
+        tzinfo=timezone.utc,
+        nanosecond_remainder=999,
+    )
+    assert Timestamp.from_datetime(dt) == ts
+
+
+def test_timestamp_dict_preserves_nanoseconds():
+    ts = Timestamp.from_dict("1970-01-01T00:00:01.123456789Z")
+
+    assert ts == Timestamp(seconds=1, nanos=123456789)
+    assert ts.to_dict() == "1970-01-01T00:00:01.123456789Z"
+
+
+def test_timestamp_dict_preserves_nanoseconds_with_offset():
+    ts = Timestamp.from_dict("1970-01-01T01:00:01.123456789+01:00")
+
+    assert ts == Timestamp(seconds=1, nanos=123456789)
+    assert ts.to_dict() == "1970-01-01T00:00:01.123456789Z"

betterproto2_compiler/src/betterproto2_compiler/known_types/__init__.py

@@ -17,8 +17,7 @@
 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.
+# The source code of the method is read from the `known_types` folder.
 KNOWN_METHODS: dict[tuple[str, str], list[Callable]] = {
     ("google.protobuf", "Any"): [Any.pack, Any.unpack, Any.to_dict, Any.from_dict],
     ("google.protobuf", "Timestamp"): [
@@ -107,6 +106,11 @@
     ],
 }
 
+# For each (package, message name), lists imports required by known-type methods.
+KNOWN_IMPORTS: dict[tuple[str, str], tuple[str, ...]] = {
+    ("google.protobuf", "Timestamp"): ("from betterproto2.nano_datetime import NanoDatetime",),
+}
+
 # A wrapped type is the type of a message that is automatically replaced by a known Python type.
 WRAPPED_TYPES: dict[tuple[str, str], str] = {
     ("google.protobuf", "BoolValue"): "bool",