diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index cbd2f27f..923c2495 100644 --- a/google/cloud/ndb/model.py +++ b/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,6 +3522,7 @@ class DateTimeProperty(Property): _auto_now = False _auto_now_add = False + _tzinfo = None def __init__( self, @@ -3521,6 +3530,7 @@ def __init__( *, 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. diff --git a/tests/system/test_crud.py b/tests/system/test_crud.py index 6108a1cb..adde1b0c 100644 --- a/tests/system/test_crud.py +++ b/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() diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index ea635573..5e4cef86 100644 --- a/tests/unit/test_model.py +++ b/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 diff --git a/tests/unit/test_query.py b/tests/unit/test_query.py index 95ba7aa2..a95c8d84 100644 --- a/tests/unit/test_query.py +++ b/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()