fix naive/aware timestamp handling in packing

fixes #34.

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]

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