Skip to content

Commit

Permalink
Prefer using built-in wrappers for timestamps & durations
Browse files Browse the repository at this point in the history
  • Loading branch information
chriso committed Jun 26, 2024
1 parent 7eebe4c commit fdbedc1
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 4 deletions.
17 changes: 17 additions & 0 deletions src/dispatch/any.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
from __future__ import annotations

import pickle
from datetime import UTC, datetime, timedelta
from typing import Any

import google.protobuf.any_pb2
import google.protobuf.duration_pb2
import google.protobuf.empty_pb2
import google.protobuf.message
import google.protobuf.timestamp_pb2
import google.protobuf.wrappers_pb2
from google.protobuf import descriptor_pool, message_factory

Expand All @@ -30,6 +33,16 @@ def marshal_any(value: Any) -> google.protobuf.any_pb2.Any:
value = google.protobuf.wrappers_pb2.StringValue(value=value)
elif isinstance(value, bytes):
value = google.protobuf.wrappers_pb2.BytesValue(value=value)
elif isinstance(value, datetime):
# Note: datetime only supports microsecond granularity
seconds = int(value.timestamp())
nanos = value.microsecond * 1000
value = google.protobuf.timestamp_pb2.Timestamp(seconds=seconds, nanos=nanos)
elif isinstance(value, timedelta):
# Note: timedelta only supports microsecond granularity
seconds = int(value.total_seconds())
nanos = value.microseconds * 1000
value = google.protobuf.duration_pb2.Duration(seconds=seconds, nanos=nanos)

if not isinstance(value, google.protobuf.message.Message):
value = pickled_pb.Pickled(pickled_value=pickle.dumps(value))
Expand Down Expand Up @@ -76,5 +89,9 @@ def unmarshal_any(any: google.protobuf.any_pb2.Any) -> Any:
except Exception as e:
# Otherwise, return the literal bytes.
return proto.value
elif isinstance(proto, google.protobuf.timestamp_pb2.Timestamp):
return proto.ToDatetime(tzinfo=UTC)
elif isinstance(proto, google.protobuf.duration_pb2.Duration):
return proto.ToTimedelta()

return proto
8 changes: 4 additions & 4 deletions tests/dispatch/test_any.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import pickle
from datetime import datetime, timedelta
from datetime import UTC, datetime, timedelta

from dispatch.any import INT64_MAX, INT64_MIN, marshal_any, unmarshal_any
from dispatch.sdk.v1 import error_pb2 as error_pb
Expand Down Expand Up @@ -70,16 +70,16 @@ def test_unmarshal_bytes():


def test_unmarshal_timestamp():
ts = datetime.fromtimestamp(
1719372909.641448
) # datetime.datetime(2024, 6, 26, 13, 35, 9, 641448)
ts = datetime.fromtimestamp(1719372909.641448, UTC)
boxed = marshal_any(ts)
assert "type.googleapis.com/google.protobuf.Timestamp" == boxed.type_url
assert ts == unmarshal_any(boxed)


def test_unmarshal_duration():
d = timedelta(seconds=1, microseconds=1234)
boxed = marshal_any(d)
assert "type.googleapis.com/google.protobuf.Duration" == boxed.type_url
assert d == unmarshal_any(boxed)


Expand Down

0 comments on commit fdbedc1

Please sign in to comment.