Skip to content

Commit

Permalink
BUG: fix pickling on Timestamp, all tests pass
Browse files Browse the repository at this point in the history
  • Loading branch information
adamklein committed Feb 22, 2012
1 parent b526c85 commit 2b2e15e
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 9 deletions.
22 changes: 16 additions & 6 deletions pandas/core/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -1150,13 +1150,13 @@ def __new__(cls, data=None,
periods=periods, offset=offset)

index = np.array(_dt_unbox_array(list(xdr)), dtype='M8[us]',
copy=False)
copy=False)

index = index.view(cls)
index.name = name
index.offset = offset
index.tzinfo = tzinfo
index.freq = freq
index.tzinfo = tzinfo

return index

Expand Down Expand Up @@ -1200,8 +1200,8 @@ def __new__(cls, data=None,
subarr = subarr.view(cls)
subarr.name = name
subarr.offset = offset
subarr.tzinfo = tzinfo
subarr.freq = freq
subarr.tzinfo = tzinfo

return subarr

Expand Down Expand Up @@ -1281,18 +1281,19 @@ def __repr__(self):
def __reduce__(self):
"""Necessary for making this object picklable"""
object_state = list(np.ndarray.__reduce__(self))
subclass_state = self.name, self.offset, self.freq
subclass_state = self.name, self.offset, self.freq, self.tzinfo
object_state[2] = (object_state[2], subclass_state)
return tuple(object_state)

def __setstate__(self, state):
"""Necessary for making this object picklable"""
if len(state) == 2:
nd_state, own_state = state
np.ndarray.__setstate__(self, nd_state)
self.name = own_state[0]
self.offset = own_state[1]
self.freq = own_state[2]
self.tzinfo = own_state[3]
np.ndarray.__setstate__(self, nd_state)
else: # pragma: no cover
np.ndarray.__setstate__(self, state)

Expand Down Expand Up @@ -1417,6 +1418,7 @@ def __array_finalize__(self, obj):
return self.item()

self.offset = getattr(obj, 'offset', None)
self.freq = getattr(obj, 'freq', None)
self.tzinfo = getattr(obj, 'tzinfo', None)

def intersection(self, other):
Expand Down Expand Up @@ -1464,14 +1466,22 @@ def __getitem__(self, key):
else:
return _dt_box(val, tzinfo=self.tzinfo)
else:
new_offset = self.offset
new_freq = self.freq
if (type(key) == slice and new_offset is not None
and key.step is not None):
new_offset = key.step * self.offset
new_freq = None

if com._is_bool_indexer(key):
key = np.asarray(key)

result = arr_idx[key]
if result.ndim > 1:
return result

return DatetimeIndex(result, name=self.name)
return DatetimeIndex(result, name=self.name, offset=new_offset,
freq=new_freq, tzinfo=self.tzinfo)

# Try to run function on index first, and then on elements of index
# Especially important for group-by functionality
Expand Down
10 changes: 10 additions & 0 deletions pandas/src/datetime.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,16 @@ class Timestamp(_Timestamp):

return ts_base

def __setstate__(self, state):
self.value = state[0]
self.offset = state[1]
self.tzinfo = state[2]

def __reduce__(self):
object_state = self.value, self.offset, self.tzinfo
return (Timestamp, object_state)


# This is PITA. Because we inherit from datetime, which has very specific
# construction requirements, we need to do object instantiation in python
# (see Timestamp class above). This will serve as a C extension type that
Expand Down
5 changes: 2 additions & 3 deletions pandas/tests/test_daterange.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,6 @@ def test_union(self):

the_union = left.union(right)
self.assert_(isinstance(the_union, Index))
self.assert_(not isinstance(the_union, DateRange))

# non-overlapping, no gap
left = self.rng[:5]
Expand All @@ -168,7 +167,7 @@ def test_union(self):
rng = DateRange(START, END, offset=datetools.bmonthEnd)

the_union = self.rng.union(rng)
self.assert_(not isinstance(the_union, DateRange))
self.assert_(isinstance(the_union, DateRange))

def test_outer_join(self):
""" should behave just as union test"""
Expand Down Expand Up @@ -228,7 +227,7 @@ def test_intersection(self):

# non-overlapping
the_int = rng[:10].intersection(rng[10:])
expected = DateRange([])
expected = DatetimeIndex([])
self.assert_(the_int.equals(expected))

def test_intersection_bug(self):
Expand Down

0 comments on commit 2b2e15e

Please sign in to comment.