Merge pull request pandas-dev#67 from manahl/multiindex-timezone-support

Add multiindex timezone support

Author: AdrianTeng

arctic/store/_pandas_ndarray_store.py

@@ -1,15 +1,18 @@
+import ast
 import logging
 
 from bson.binary import Binary
 from pandas import DataFrame, MultiIndex, Series, DatetimeIndex, Panel
+from pandas.tseries.index import DatetimeIndex
 from pandas.tslib import Timestamp, get_timezone
+
 import numpy as np
-import ast
 
 from .._compression import compress, decompress
+from ..date._util import to_pandas_closed_closed
 from ..exceptions import ArcticException
 from ._ndarray_store import NdarrayStore
-from ..date._util import to_pandas_closed_closed
+
 
 log = logging.getLogger(__name__)
 
@@ -57,6 +60,8 @@ def _index_to_records(self, df):
 
         if isinstance(index, DatetimeIndex) and index.tz is not None:
             metadata['index_tz'] = get_timezone(index.tz)
+        elif isinstance(index, MultiIndex):
+            metadata['index_tz'] = [get_timezone(i.tz) if isinstance(i, DatetimeIndex) else None for i in index.levels]
 
         return index_names, ix_vals, metadata
 
@@ -66,6 +71,10 @@ def _index_from_records(self, recarr):
 
         if isinstance(rtn, DatetimeIndex) and 'index_tz' in recarr.dtype.metadata:
             rtn = rtn.tz_localize('UTC').tz_convert(recarr.dtype.metadata['index_tz'])
+        elif isinstance(rtn, MultiIndex):
+            for i, tz in enumerate(recarr.dtype.metadata.get('index_tz')):
+                if tz is not None:
+                    rtn.set_levels(rtn.levels[i].tz_localize('UTC').tz_convert(tz), i, inplace=True)
 
         return rtn
 

tests/integration/store/test_bitemporal_store.py

@@ -28,6 +28,8 @@
                                 2012-11-08 17:06:11.040 |  3.0
                                 2012-11-09 17:06:11.040 |  3.5""")
 
+LOCAL_TZ = mktz()
+
 
 def test_new_ts_read_write(bitemporal_library):
     bitemporal_library.update('spam', ts1)
@@ -57,7 +59,7 @@ def test_write_ts_with_column_name_same_as_observed_dt_ok(bitemporal_library):
 def test_last_update(bitemporal_library):
     bitemporal_library.update('spam', ts1, as_of=dt(2015, 1, 1))
     bitemporal_library.update('spam', ts1, as_of=dt(2015, 1, 2))
-    assert bitemporal_library.read('spam').last_updated == dt(2015, 1, 2)
+    assert bitemporal_library.read('spam').last_updated == dt(2015, 1, 2, tzinfo=LOCAL_TZ)
 
 
 def test_existing_ts_update_and_read(bitemporal_library):
@@ -195,10 +197,9 @@ def test_read_ts_raw_all_version_ok(bitemporal_library):
 
 
 def test_bitemporal_store_saves_as_of_with_timezone(bitemporal_library):
-    local_tz = mktz()
     bitemporal_library.update('spam', ts1, as_of=dt(2015, 5, 1))
     df = bitemporal_library.read('spam', raw=True).data
-    assert all([x[1].replace(tzinfo=mktz('UTC')) == dt(2015, 5, 1, tzinfo=local_tz) for x in df.index])
+    assert all([x[1] == dt(2015, 5, 1, tzinfo=LOCAL_TZ) for x in df.index])
 
 
 def test_bitemporal_store_read_as_of_timezone(bitemporal_library):
@@ -263,4 +264,4 @@ def test_multi_index_update(bitemporal_library):
     bitemporal_library.update('spam', ts, as_of=dt(2015, 1, 1))
     bitemporal_library.update('spam', ts2, as_of=dt(2015, 1, 2))
     assert_frame_equal(expected_ts, bitemporal_library.read('spam').data)
-    assert bitemporal_library.read('spam').last_updated == dt(2015, 1, 2)
+    assert bitemporal_library.read('spam').last_updated == dt(2015, 1, 2, tzinfo=LOCAL_TZ)

tests/integration/store/test_pandas_store.py

@@ -1,20 +1,22 @@
 from StringIO import StringIO
 from datetime import datetime as dt, timedelta as dtd
+import io
+import itertools
+import string
+
 from dateutil.rrule import rrule, DAILY
+from mock import Mock, patch
 from pandas import DataFrame, Series, DatetimeIndex, MultiIndex, read_csv, Panel, date_range, concat
+from pandas.tseries.offsets import DateOffset
 from pandas.util.testing import assert_frame_equal, assert_series_equal
-import numpy as np
 import pytest
-import io
-import itertools
-from mock import Mock, patch
-import string
 
-from arctic.date import DateRange, mktz
 from arctic._compression import decompress
+from arctic.date import DateRange, mktz
 from arctic.store._pandas_ndarray_store import PandasDataFrameStore, PandasSeriesStore, PandasStore
 from arctic.store.version_store import register_versioned_storage
-from pandas.tseries.offsets import DateOffset
+import numpy as np
+from tests.util import read_str_as_pandas
 
 
 register_versioned_storage(PandasDataFrameStore)
@@ -831,9 +833,21 @@ def test_data_info_cols(library):
     library.write('test_data', s)
     md = library.get_info('test_data')
     assert md == {'dtype': [('level_0', '<i8'), ('level_1', 'S2'), ('0', '<i8')],
-                  'col_names': {u'index': [u'level_0', u'level_1'], u'columns': [u'0']},
+                  'col_names': {u'index': [u'level_0', u'level_1'], u'columns': [u'0'], u'index_tz': [None, None]},
                   'type': u'pandasdf',
                   'handler': 'PandasDataFrameStore',
                   'rows': 3,
                   'segment_count': 1,
                   'size': 54}
+
+
+def test_read_write_multiindex_store_keeps_timezone(library):
+    """If I write a multi-index dataframe and reads it, the timezone of the index shouldn't change, right?"""
+    hk, ny, ldn = mktz('Asia/Hong_Kong'), mktz('America/New_York'), mktz('Europe/London')
+    row0 = [dt(2015, 1, 1, tzinfo=hk), dt(2015, 1, 1, tzinfo=ny), dt(2015, 1, 1, tzinfo=ldn), 0, 42]
+    row1 = [dt(2015, 1, 2, tzinfo=hk), dt(2015, 1, 2, tzinfo=ny), dt(2015, 1, 2, tzinfo=ldn), 1, 43]
+    df = DataFrame([row0, row1], columns=['dt_a', 'dt_b', 'dt_c', 'index_0', 'data'])
+    df = df.set_index(['dt_a', 'dt_b', 'dt_c', 'index_0'])
+    library.write('spam', df)
+    assert list(library.read('spam').data.index[0]) == row0[:-1]
+    assert list(library.read('spam').data.index[1]) == row1[:-1]