diff --git a/test_umsgpack.py b/test_umsgpack.py index eeea33e..42a6748 100644 --- a/test_umsgpack.py +++ b/test_umsgpack.py @@ -319,6 +319,18 @@ ["float precision double", 2.5, b"\xcb\x40\x04\x00\x00\x00\x00\x00\x00"], ] +naive_timestamp_test_vectors = [ + ["32-bit timestamp (naive)", datetime.datetime(2000, 1, 1, 10, 5, 2, 0, umsgpack._utc_tzinfo), + b"\xd6\xff\x38\x6d\xd1\x4e", + datetime.datetime(2000, 1, 1, 10, 5, 2, 0, umsgpack._utc_tzinfo)], + ["64-bit timestamp (naive)", datetime.datetime(2200, 1, 1, 10, 5, 2, 1234, umsgpack._utc_tzinfo), + b"\xd7\xff\x00\x4b\x51\x41\xb0\x9e\xa6\xce", + datetime.datetime(2200, 1, 1, 10, 5, 2, 1234, umsgpack._utc_tzinfo)], + ["96-bit timestamp (naive)", datetime.datetime(3000, 1, 1, 10, 5, 2, 1234, umsgpack._utc_tzinfo), + b"\xc7\x0c\xff\x00\x12\xd4\x50\x00\x00\x00\x07\x91\x5f\x59\xce", + datetime.datetime(3000, 1, 1, 10, 5, 2, 1234, umsgpack._utc_tzinfo)], +] + CustomType = namedtuple('CustomType', ['x', 'y', 'z']) ext_handlers = { @@ -540,6 +552,24 @@ def test_pack_force_float_precision(self): packed = umsgpack.packb(obj, force_float_precision=precision) self.assertEqual(packed, data) + def test_pack_naive_timestamp(self): + for (name, obj, data, _) in naive_timestamp_test_vectors: + obj_repr = repr(obj) + print("\t Testing %s: object %s" % + (name, obj_repr if len(obj_repr) < 24 else obj_repr[0:24] + "...")) + + packed = umsgpack.packb(obj) + self.assertEqual(packed, data) + + def test_unpack_naive_timestamp(self): + for (name, _, data, obj) in naive_timestamp_test_vectors: + obj_repr = repr(obj) + print("\t Testing %s: object %s" % + (name, obj_repr if len(obj_repr) < 24 else obj_repr[0:24] + "...")) + + unpacked = umsgpack.unpackb(data) + self.assertEqual(unpacked, obj) + def test_pack_ext_override(self): # Test overridden packing of datetime.datetime (name, obj, data) = override_ext_handlers_test_vectors[0] diff --git a/umsgpack.py b/umsgpack.py index 372455b..bb0d6c8 100644 --- a/umsgpack.py +++ b/umsgpack.py @@ -348,7 +348,14 @@ def _pack_ext(obj, fp, options): def _pack_ext_timestamp(obj, fp, options): - delta = obj - _epoch + if not obj.tzinfo: + # Object is naive datetime, convert to aware date time, + # assuming UTC timezone + delta = obj.replace(tzinfo=_utc_tzinfo) - _epoch + else: + # Object is aware datetime + delta = obj - _epoch + seconds = delta.seconds + delta.days * 86400 microseconds = delta.microseconds @@ -1052,9 +1059,21 @@ def __init(): if sys.version_info[0] == 3: _utc_tzinfo = datetime.timezone.utc else: - _utc_tzinfo = None + class UTC(datetime.tzinfo): + ZERO = datetime.timedelta(0) + + def utcoffset(self, dt): + return UTC.ZERO + + def tzname(self, dt): + return "UTC" + + def dst(self, dt): + return UTC.ZERO + + _utc_tzinfo = UTC() - # Calculate epoch datetime + # Calculate an aware epoch datetime _epoch = datetime.datetime(1970, 1, 1, tzinfo=_utc_tzinfo) # Auto-detect system float precision