Add tzinfo to DateTimeProperty. (#226)

Fixes #7.

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/model.py

@@ -254,6 +254,8 @@ class Person(Model):
 import six
 import zlib
 
+import pytz
+
 from google.cloud.datastore import entity as ds_entity_module
 from google.cloud.datastore import helpers
 from google.cloud.datastore_v1.proto import entity_pb2
@@ -3467,9 +3469,12 @@ def _validate(self, value):
 class DateTimeProperty(Property):
     """A property that contains :class:`~datetime.datetime` values.
 
-    This property expects "naive" datetime stamps, i.e. no timezone can
-    be set. Furthermore, the assumption is that naive datetime stamps
-    represent UTC.
+    If ``tzinfo`` is not set, this property expects "naive" datetime stamps,
+    i.e. no timezone can be set. Furthermore, the assumption is that naive
+    datetime stamps represent UTC.
+
+    If ``tzinfo`` is set, timestamps will be stored as UTC and converted back
+    to the timezone set by ``tzinfo`` when reading values back out.
 
     .. note::
 
@@ -3493,6 +3498,9 @@ class DateTimeProperty(Property):
             updated.
         auto_now_add (bool): Indicates that the property should be set to the
             current datetime when an entity is created.
+        tzinfo (Optional[datetime.tzinfo]): If set, values read from Datastore
+            will be converted to this timezone. Otherwise, values will be
+            returned as naive datetime objects with an implied UTC timezone.
         indexed (bool): Indicates if the value should be indexed.
         repeated (bool): Indicates if this property is repeated, i.e. contains
             multiple values.
@@ -3514,13 +3522,15 @@ class DateTimeProperty(Property):
 
     _auto_now = False
     _auto_now_add = False
+    _tzinfo = None
 
     def __init__(
         self,
         name=None,
         *,
         auto_now=None,
         auto_now_add=None,
+        tzinfo=None,
         indexed=None,
         repeated=None,
         required=None,
@@ -3556,6 +3566,8 @@ def __init__(
             self._auto_now = auto_now
         if auto_now_add is not None:
             self._auto_now_add = auto_now_add
+        if tzinfo is not None:
+            self._tzinfo = tzinfo
 
     def _validate(self, value):
         """Validate a ``value`` before setting it.
@@ -3571,10 +3583,10 @@ def _validate(self, value):
                 "Expected datetime, got {!r}".format(value)
             )
 
-        if value.tzinfo is not None:
+        if self._tzinfo is None and value.tzinfo is not None:
             raise exceptions.BadValueError(
-                "DatetimeProperty {} can only support naive datetimes "
-                "(presumed UTC). Please derive a new Property to support "
+                "DatetimeProperty without tzinfo {} can only support naive "
+                "datetimes (presumed UTC). Please set tzinfo to support "
                 "alternate timezones.".format(self._name)
             )
 
@@ -3613,12 +3625,32 @@ def _from_base_type(self, value):
             value (datetime.datetime): The value to be converted.
 
         Returns:
-            Optional[datetime.datetime]: The value without ``tzinfo`` or
-                ``None`` if value did not have ``tzinfo`` set.
+            Optional[datetime.datetime]: If ``tzinfo`` is set on this property,
+                the value converted to the timezone in ``tzinfo``. Otherwise
+                returns the value without ``tzinfo`` or ``None`` if value did
+                not have ``tzinfo`` set.
         """
-        if value.tzinfo is not None:
+        if self._tzinfo is not None:
+            return value.astimezone(self._tzinfo)
+
+        elif value.tzinfo is not None:
             return value.replace(tzinfo=None)
 
+    def _to_base_type(self, value):
+        """Convert a value to the "base" value type for this property.
+
+        Args:
+            value (datetime.datetime): The value to be converted.
+
+        Returns:
+            google.cloud.datastore.Key: The converted value.
+
+        Raises:
+            TypeError: If ``value`` is not a :class:`~key.Key`.
+        """
+        if self._tzinfo is not None and value.tzinfo is not None:
+            return value.astimezone(pytz.utc)
+
 
 class DateProperty(DateTimeProperty):
     """A property that contains :class:`~datetime.date` values.

packages/google-cloud-ndb/tests/system/test_crud.py

@@ -261,6 +261,37 @@ class SomeKind(ndb.Model):
     dispose_of(key._key)
 
 
+@pytest.mark.usefixtures("client_context")
+def test_datetime_w_tzinfo(dispose_of, ds_client):
+    class timezone(datetime.tzinfo):
+        def __init__(self, offset):
+            self.offset = datetime.timedelta(hours=offset)
+
+        def utcoffset(self, dt):
+            return self.offset
+
+        def dst(self, dt):
+            return datetime.timedelta(0)
+
+    mytz = timezone(-4)
+
+    class SomeKind(ndb.Model):
+        foo = ndb.DateTimeProperty(tzinfo=mytz)
+        bar = ndb.DateTimeProperty(tzinfo=mytz)
+
+    entity = SomeKind(
+        foo=datetime.datetime(2010, 5, 12, 2, 42, tzinfo=timezone(-5)),
+        bar=datetime.datetime(2010, 5, 12, 2, 42),
+    )
+    key = entity.put()
+
+    retrieved = key.get()
+    assert retrieved.foo == datetime.datetime(2010, 5, 12, 3, 42, tzinfo=mytz)
+    assert retrieved.bar == datetime.datetime(2010, 5, 11, 22, 42, tzinfo=mytz)
+
+    dispose_of(key._key)
+
+
 def test_parallel_threads(dispose_of, namespace):
     client = ndb.Client(namespace=namespace)
 
@@ -337,7 +368,7 @@ class SomeKind(ndb.Model):
 
 @pytest.mark.usefixtures("client_context")
 def test_retrieve_entity_with_legacy_compressed_property(
-    ds_entity_with_meanings
+    ds_entity_with_meanings,
 ):
     class SomeKind(ndb.Model):
         blob = ndb.BlobProperty()

packages/google-cloud-ndb/tests/unit/test_model.py

@@ -38,6 +38,20 @@
 from tests.unit import utils
 
 
+class timezone(datetime.tzinfo):
+    def __init__(self, offset):
+        self.offset = datetime.timedelta(hours=offset)
+
+    def utcoffset(self, dt):
+        return self.offset
+
+    def dst(self, dt):
+        return datetime.timedelta(0)
+
+    def __eq__(self, other):
+        return self.offset == other.offset
+
+
 def test___all__():
     utils.verify___all__(model)
 
@@ -2548,6 +2562,7 @@ def test_constructor_explicit():
             name="dt_val",
             auto_now=True,
             auto_now_add=False,
+            tzinfo=timezone(-4),
             indexed=False,
             repeated=False,
             required=True,
@@ -2559,6 +2574,7 @@ def test_constructor_explicit():
         assert prop._name == "dt_val"
         assert prop._auto_now
         assert not prop._auto_now_add
+        assert prop._tzinfo == timezone(-4)
         assert not prop._indexed
         assert not prop._repeated
         assert prop._required
@@ -2671,6 +2687,28 @@ def test__from_base_type_timezone():
         value = datetime.datetime(2010, 5, 12, tzinfo=pytz.utc)
         assert prop._from_base_type(value) == datetime.datetime(2010, 5, 12)
 
+    @staticmethod
+    def test__from_base_type_convert_timezone():
+        prop = model.DateTimeProperty(name="dt_val", tzinfo=timezone(-4))
+        value = datetime.datetime(2010, 5, 12, tzinfo=pytz.utc)
+        assert prop._from_base_type(value) == datetime.datetime(
+            2010, 5, 11, 20, tzinfo=timezone(-4)
+        )
+
+    @staticmethod
+    def test__to_base_type_noop():
+        prop = model.DateTimeProperty(name="dt_val", tzinfo=timezone(-4))
+        value = datetime.datetime(2010, 5, 12)
+        assert prop._to_base_type(value) is None
+
+    @staticmethod
+    def test__to_base_type_convert_to_utc():
+        prop = model.DateTimeProperty(name="dt_val", tzinfo=timezone(-4))
+        value = datetime.datetime(2010, 5, 12, tzinfo=timezone(-4))
+        assert prop._to_base_type(value) == datetime.datetime(
+            2010, 5, 12, 4, tzinfo=pytz.utc
+        )
+
 
 class TestDateProperty:
     @staticmethod

packages/google-cloud-ndb/tests/unit/test_query.py

@@ -1185,7 +1185,7 @@ class Bar(model.Model):
     @pytest.mark.usefixtures("in_context")
     @unittest.mock.patch("google.cloud.ndb.query._datastore_query")
     def test_constructor_with_class_attribute_projection_and_distinct(
-        _datastore_query
+        _datastore_query,
     ):
         class Foo(model.Model):
             string_attr = model.StringProperty()