Calculate unix timestamp using timezone if present

riak/tests/test_datetime.py

@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+import datetime
+import unittest
+
+from riak.util import epoch, epoch_tz, \
+        unix_time_millis
+
+# NB: without tzinfo, this is UTC
+ts0 = datetime.datetime(2015, 1, 1, 12, 1, 2, 987000)
+ts0_ts = 1420113662987
+ts0_ts_pst = 1420142462987
+
+
+class DatetimeUnitTests(unittest.TestCase):
+    def test_get_unix_time_without_tzinfo(self):
+        self.assertIsNone(epoch.tzinfo)
+        self.assertIsNotNone(epoch_tz.tzinfo)
+        self.assertIsNone(ts0.tzinfo)
+        utm = unix_time_millis(ts0)
+        self.assertEqual(utm, ts0_ts)
+
+    def test_get_unix_time_with_tzinfo(self):
+        try:
+            import pytz
+            tz = pytz.timezone('America/Los_Angeles')
+            ts0_pst = tz.localize(ts0)
+            utm = unix_time_millis(ts0_pst)
+            self.assertEqual(utm, ts0_ts_pst)
+        except ImportError:
+            pass

riak/tests/test_timeseries_pbuf.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 import datetime
+import six
 import unittest
 
 import riak.pb.riak_ts_pb2
@@ -457,5 +458,6 @@ def test_store_and_fetch_gh_483(self):
 
         row = ts_obj.rows[0]
         self.assertEqual(len(row), 5)
-        exp = rows[0]
+        exp = [six.b('hash1'), six.b('user2'), now,
+               six.b('frazzle'), 12.3]
         self.assertEqual(row, exp)

riak/tz.py

@@ -0,0 +1,18 @@
+from datetime import tzinfo, timedelta
+
+ZERO = timedelta(0)
+
+
+class UTC(tzinfo):
+    """UTC"""
+
+    def utcoffset(self, dt):
+        return ZERO
+
+    def tzname(self, dt):
+        return "UTC"
+
+    def dst(self, dt):
+        return ZERO
+
+utc = UTC()

riak/util.py

@@ -8,10 +8,19 @@
 from six import string_types, PY2
 
 epoch = datetime.datetime.utcfromtimestamp(0)
+try:
+    import pytz
+    epoch_tz = pytz.utc.localize(epoch)
+except ImportError:
+    from riak.tz import utc
+    epoch_tz = datetime.datetime.fromtimestamp(0, tz=utc)
 
 
 def unix_time_millis(dt):
-    td = dt - epoch
+    if dt.tzinfo:
+        td = dt - epoch_tz
+    else:
+        td = dt - epoch
     tdms = ((td.days * 24 * 3600) + td.seconds) * 1000
     ms = td.microseconds // 1000
     return tdms + ms

tox.ini

@@ -12,5 +12,7 @@ basepython = {env:HOME}/.pyenv/versions/riak-py278/bin/python2.7
 [testenv]
 install_command = pip install --upgrade {packages}
 commands = {envpython} setup.py test
-deps = pip
+deps =
+    pip
+    pytz
 passenv = RUN_* SKIP_* RIAK_*