Fix dict encoding for timezone aware datetimes (#468)

Co-authored-by: James Hilton-Balfe <[email protected]>

Author: semicolonTransistor

src/betterproto/__init__.py

@@ -1583,6 +1583,9 @@ def to_datetime(self) -> datetime:
     @staticmethod
     def timestamp_to_json(dt: datetime) -> str:
         nanos = dt.microsecond * 1e3
+        if dt.tzinfo is not None:
+            # change timezone aware datetime objects to utc
+            dt = dt.astimezone(timezone.utc)
         copy = dt.replace(microsecond=0, tzinfo=None)
         result = copy.isoformat()
         if (nanos % 1e9) == 0:

tests/inputs/timestamp_dict_encode/test_timestamp_dict_encode.py

@@ -0,0 +1,82 @@
+from datetime import (
+    datetime,
+    timedelta,
+    timezone,
+)
+
+import pytest
+
+from tests.output_betterproto.timestamp_dict_encode import Test
+
+
+# Current World Timezone range (UTC-12 to UTC+14)
+MIN_UTC_OFFSET_MIN = -12 * 60
+MAX_UTC_OFFSET_MIN = 14 * 60
+
+# Generate all timezones in range in 15 min increments
+timezones = [
+    timezone(timedelta(minutes=x))
+    for x in range(MIN_UTC_OFFSET_MIN, MAX_UTC_OFFSET_MIN + 1, 15)
+]
+
+
+@pytest.mark.parametrize("tz", timezones)
+def test_timezone_aware_datetime_dict_encode(tz: timezone):
+    original_time = datetime.now(tz=tz)
+    original_message = Test()
+    original_message.ts = original_time
+    encoded = original_message.to_dict()
+    decoded_message = Test()
+    decoded_message.from_dict(encoded)
+
+    # check that the timestamps are equal after decoding from dict
+    assert original_message.ts.tzinfo is not None
+    assert decoded_message.ts.tzinfo is not None
+    assert original_message.ts == decoded_message.ts
+
+
+def test_naive_datetime_dict_encode():
+    # make suer naive datetime objects are still treated as utc
+    original_time = datetime.now()
+    assert original_time.tzinfo is None
+    original_message = Test()
+    original_message.ts = original_time
+    original_time_utc = original_time.replace(tzinfo=timezone.utc)
+    encoded = original_message.to_dict()
+    decoded_message = Test()
+    decoded_message.from_dict(encoded)
+
+    # check that the timestamps are equal after decoding from dict
+    assert decoded_message.ts.tzinfo is not None
+    assert original_time_utc == decoded_message.ts
+
+
+@pytest.mark.parametrize("tz", timezones)
+def test_timezone_aware_json_serialize(tz: timezone):
+    original_time = datetime.now(tz=tz)
+    original_message = Test()
+    original_message.ts = original_time
+    json_serialized = original_message.to_json()
+    decoded_message = Test()
+    decoded_message.from_json(json_serialized)
+
+    # check that the timestamps are equal after decoding from dict
+    assert original_message.ts.tzinfo is not None
+    assert decoded_message.ts.tzinfo is not None
+    assert original_message.ts == decoded_message.ts
+
+
+def test_naive_datetime_json_serialize():
+    # make suer naive datetime objects are still treated as utc
+    original_time = datetime.now()
+    assert original_time.tzinfo is None
+    original_message = Test()
+    original_message.ts = original_time
+    original_time_utc = original_time.replace(tzinfo=timezone.utc)
+    json_serialized = original_message.to_json()
+    decoded_message = Test()
+    decoded_message.from_json(json_serialized)
+
+    # check that the timestamps are equal after decoding from dict
+    assert decoded_message.ts.tzinfo is not None
+    assert original_time_utc == decoded_message.ts

tests/inputs/timestamp_dict_encode/timestamp_dict_encode.json

@@ -0,0 +1,3 @@
+{
+    "ts" : "2023-03-15T22:35:51.253277Z"
+}

tests/inputs/timestamp_dict_encode/timestamp_dict_encode.proto

@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+package timestamp_dict_encode;
+
+import "google/protobuf/timestamp.proto";
+
+message Test {
+  google.protobuf.Timestamp ts = 1;
+}