ENH: store pytz time zones as zone strings in HDFStore, close #1232

Author: wesm

pandas/io/pytables.py

@@ -585,6 +585,9 @@ def _write_index(self, group, key, index):
             if hasattr(index, 'freq'):
                 node._v_attrs.freq = index.freq
 
+            if hasattr(index, 'tz') and index.tz is not None:
+                node._v_attrs.tz = index.tz.zone
+
     def _read_index(self, group, key):
         variety = getattr(group._v_attrs, '%s_variety' % key)
 
@@ -668,15 +671,21 @@ def _read_index_node(self, node):
             name = node._v_attrs.name
 
         index_class = getattr(node._v_attrs, 'index_class', Index)
+
+        factory = _get_index_factory(index_class)
+
         kwargs = {}
         if 'freq' in node._v_attrs:
             kwargs['freq'] = node._v_attrs['freq']
 
+        if 'tz' in node._v_attrs:
+            kwargs['tz'] = node._v_attrs['tz']
+
         if kind in ('date', 'datetime'):
-            index = index_class(_unconvert_index(data, kind), dtype=object,
-                                **kwargs)
+            index = factory(_unconvert_index(data, kind), dtype=object,
+                            **kwargs)
         else:
-            index = index_class(_unconvert_index(data, kind), **kwargs)
+            index = factory(_unconvert_index(data, kind), **kwargs)
 
         index.name = name
 
@@ -1085,3 +1094,11 @@ def select_coords(self):
         """
         self.values = self.table.getWhereList(self.the_condition)
 
+def _get_index_factory(klass):
+    if klass == DatetimeIndex:
+        def f(values, freq=None, tz=None):
+            return DatetimeIndex._simple_new(values, None, freq=freq,
+                                             tz=tz)
+        return f
+    return klass
+

pandas/io/tests/test_pytables.py

@@ -8,7 +8,8 @@
 from datetime import datetime
 import numpy as np
 
-from pandas import Series, DataFrame, Panel, MultiIndex, bdate_range
+from pandas import (Series, DataFrame, Panel, MultiIndex, bdate_range,
+                    date_range)
 from pandas.io.pytables import HDFStore, get_store
 import pandas.util.testing as tm
 from pandas.tests.test_series import assert_series_equal
@@ -338,6 +339,19 @@ def test_can_serialize_dates(self):
         frame = DataFrame(np.random.randn(len(rng), 4), index=rng)
         self._check_roundtrip(frame, tm.assert_frame_equal)
 
+    def test_timezones(self):
+        rng = date_range('1/1/2000', '1/30/2000', tz='US/Eastern')
+        frame = DataFrame(np.random.randn(len(rng), 4), index=rng)
+        try:
+            store = HDFStore(self.scratchpath)
+            store['frame'] = frame
+            recons = store['frame']
+            self.assert_(recons.index.equals(rng))
+            self.assertEquals(rng.tz, recons.index.tz)
+        finally:
+            store.close()
+            os.remove(self.scratchpath)
+
     def test_store_hierarchical(self):
         index = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux'],
                                    ['one', 'two', 'three']],

pandas/tseries/index.py

@@ -322,11 +322,11 @@ def _generate(cls, start, end, periods, name, offset,
         return index
 
     @classmethod
-    def _simple_new(cls, values, name, offset, tz):
+    def _simple_new(cls, values, name, freq=None, tz=None):
         result = values.view(cls)
         result.name = name
-        result.offset = offset
-        result.tz = tz
+        result.offset = freq
+        result.tz = tools._maybe_get_tz(tz)
 
         return result