Skip to content

Commit

Permalink
ENH: store pytz time zones as zone strings in HDFStore, close #1232
Browse files Browse the repository at this point in the history
  • Loading branch information
wesm committed May 14, 2012
1 parent 7ac1b51 commit 2393ba9
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 7 deletions.
23 changes: 20 additions & 3 deletions pandas/io/pytables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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

16 changes: 15 additions & 1 deletion pandas/io/tests/test_pytables.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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']],
Expand Down
6 changes: 3 additions & 3 deletions pandas/tseries/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit 2393ba9

Please sign in to comment.