Merge pull request #8370 from jreback/tz_pickle

BUG/COMPAT: set tz on DatetimeIndex on pickle reconstruction (GH8367)

Author: jreback

doc/source/v0.15.0.txt

@@ -534,7 +534,7 @@ Internal Refactoring
 
 In 0.15.0 ``Index`` has internally been refactored to no longer sub-class ``ndarray``
 but instead subclass ``PandasObject``, similarly to the rest of the pandas objects. This change allows very easy sub-classing and creation of new index types. This should be
-a transparent change with only very limited API implications (:issue:`5080`, :issue:`7439`, :issue:`7796`, :issue:`8024`)
+a transparent change with only very limited API implications (:issue:`5080`, :issue:`7439`, :issue:`7796`, :issue:`8024`, :issue:`8367`)
 
 - you may need to unpickle pandas version < 0.15.0 pickles using ``pd.read_pickle`` rather than ``pickle.load``. See :ref:`pickle docs <io.pickle>`
 - when plotting with a ``PeriodIndex``. The ``matplotlib`` internal axes will now be arrays of ``Period`` rather than a ``PeriodIndex``. (this is similar to how a ``DatetimeIndex`` passes arrays of ``datetimes`` now)

pandas/core/index.py

@@ -3323,8 +3323,8 @@ def __contains__(self, key):
 
     def __reduce__(self):
         """Necessary for making this object picklable"""
-        d = dict(levels = [lev.view(np.ndarray) for lev in self.levels],
-                 labels = [label.view(np.ndarray) for label in self.labels],
+        d = dict(levels = [lev for lev in self.levels],
+                 labels = [label for label in self.labels],
                  sortorder = self.sortorder,
                  names = list(self.names))
         return _new_Index, (self.__class__, d), None

pandas/tests/test_index.py

@@ -1645,6 +1645,14 @@ def test_numeric_compat(self):
                       lambda : pd.date_range('2000-01-01', periods=3) * np.timedelta64(1, 'D').astype('m8[ns]') ]:
                 self.assertRaises(TypeError, f)
 
+    def test_roundtrip_pickle_with_tz(self):
+
+        # GH 8367
+        # round-trip of timezone
+        index=date_range('20130101',periods=3,tz='US/Eastern',name='foo')
+        unpickled = self.round_trip_pickle(index)
+        self.assertTrue(index.equals(unpickled))
+
 class TestPeriodIndex(Base, tm.TestCase):
     _holder = PeriodIndex
     _multiprocess_can_split_ = True
@@ -2347,6 +2355,14 @@ def test_legacy_v2_unpickle(self):
         assert_almost_equal(res, exp)
         assert_almost_equal(exp, exp2)
 
+    def test_roundtrip_pickle_with_tz(self):
+
+        # GH 8367
+        # round-trip of timezone
+        index=MultiIndex.from_product([[1,2],['a','b'],date_range('20130101',periods=3,tz='US/Eastern')],names=['one','two','three'])
+        unpickled = self.round_trip_pickle(index)
+        self.assertTrue(index.equal_levels(unpickled))
+
     def test_from_tuples_index_values(self):
         result = MultiIndex.from_tuples(self.index)
         self.assertTrue((result.values == self.index.values).all())

pandas/tseries/index.py

@@ -105,6 +105,17 @@ def _ensure_datetime64(other):
 
 _midnight = time(0, 0)
 
+def _new_DatetimeIndex(cls, d):
+    """ This is called upon unpickling, rather than the default which doesn't have arguments
+        and breaks __new__ """
+
+    # simply set the tz
+    # data are already in UTC
+    tz = d.pop('tz',None)
+    result = cls.__new__(cls, **d)
+    result.tz = tz
+    return result
+
 class DatetimeIndex(DatetimeIndexOpsMixin, Int64Index):
     """
     Immutable ndarray of datetime64 data, represented internally as int64, and
@@ -583,6 +594,15 @@ def _formatter_func(self):
         formatter = _get_format_datetime64(is_dates_only=self._is_dates_only)
         return lambda x: formatter(x, tz=self.tz)
 
+    def __reduce__(self):
+
+        # we use a special reudce here because we need
+        # to simply set the .tz (and not reinterpret it)
+
+        d = dict(data=self._data)
+        d.update(self._get_attributes_dict())
+        return _new_DatetimeIndex, (self.__class__, d), None
+
     def __setstate__(self, state):
         """Necessary for making this object picklable"""
         if isinstance(state, dict):