TickStore.max_date: convert to local timezone

Author: mckelvin

.gitignore

@@ -1,4 +1,5 @@
 /build
+*.eggs/
 *.pyc
 /src/_compress.c
 *.egg
@@ -13,3 +14,5 @@ htmlcov
 coverage.xml
 junit.xml
 /tmp/
+
+*.sw[op]

arctic/date/__init__.py

@@ -1,5 +1,5 @@
 from ._daterange import DateRange
 from ._generalslice import OPEN_CLOSED, CLOSED_OPEN, OPEN_OPEN, CLOSED_CLOSED
 from ._util import datetime_to_ms, ms_to_datetime
-from ._util import string_to_daterange, to_pandas_closed_closed, to_dt
+from ._util import string_to_daterange, to_pandas_closed_closed, to_dt, utc_dt_to_local_dt
 from ._mktz import mktz, TimezoneError

arctic/date/_util.py

@@ -166,3 +166,18 @@ def datetime_to_ms(d):
         return long((calendar.timegm(_add_tzone(d).utctimetuple()) + d.microsecond / 1000000.0) * 1e3)
     except AttributeError:
         raise TypeError('expect Python datetime object, not %s' % type(d))
+
+
+def utc_dt_to_local_dt(dtm):
+    """Convert a UTC datetime to datetime in local timezone"""
+    utc_zone = mktz("UTC")
+    if dtm.tzinfo is not None and dtm.tzinfo != utc_zone:
+        raise ValueError(
+            "Expected dtm without tzinfo or with UTC, not %r" % (
+                dtm.tzinfo
+            )
+        )
+
+    if dtm.tzinfo is None:
+        dtm = dtm.replace(tzinfo=utc_zone)
+    return dtm.astimezone(mktz())

arctic/tickstore/tickstore.py

@@ -13,7 +13,7 @@
 from pymongo.errors import OperationFailure
 from six import iteritems, string_types
 
-from ..date import DateRange, to_pandas_closed_closed, mktz, datetime_to_ms, ms_to_datetime, CLOSED_CLOSED, to_dt
+from ..date import DateRange, to_pandas_closed_closed, mktz, datetime_to_ms, ms_to_datetime, CLOSED_CLOSED, to_dt, utc_dt_to_local_dt
 from ..decorators import mongo_retry
 from ..exceptions import OverlappingDataException, NoDataFoundException, UnorderedDataException, UnhandledDtypeException, ArcticException
 from .._util import indent
@@ -699,4 +699,4 @@ def max_date(self, symbol):
         """
         res = self._collection.find_one({SYMBOL: symbol}, projection={ID: 0, END: 1},
                                         sort=[(START, pymongo.DESCENDING)])
-        return res[END]
+        return utc_dt_to_local_dt(res[END])

tests/unit/date/test_util.py

@@ -1,9 +1,10 @@
 import pytest
 import pytz
+from mock import patch
 
 from datetime import datetime as dt
 from arctic.date import datetime_to_ms, ms_to_datetime, mktz, to_pandas_closed_closed, DateRange, OPEN_OPEN, CLOSED_CLOSED
-from arctic.date._util import to_dt
+from arctic.date._util import to_dt, utc_dt_to_local_dt
 
 
 @pytest.mark.parametrize('pdt', [
@@ -105,3 +106,15 @@ def test_daterange_lt():
     assert(dr2 < dr)
     dr.start = None
     assert((dr2 < dr) == False)
+
+
+@patch("arctic.date._util.mktz", lambda zone="Asia/Shanghai": mktz(zone))
+def test_utc_dt_to_local_dt():
+    with pytest.raises(ValueError):
+        assert(utc_dt_to_local_dt(
+            dt(2000, 1, 1, 0, 0, 0, tzinfo=mktz("Asia/Shanghai"))
+        ))
+
+    utc_time = dt(2000, 1, 1, 10, 0, 0)
+    pek_time = utc_dt_to_local_dt(utc_time)  # GMT +0800
+    assert(pek_time.hour - utc_time.hour == 8)