Merge pull request #4791 from jreback/pickle_fix

BUG: pickle failing on FrozenList, when using MultiIndex (GH4788)

Author: jreback

pandas/core/base.py

@@ -108,6 +108,9 @@ def __mul__(self, other):
 
     __imul__ = __mul__
 
+    def __reduce__(self):
+        return self.__class__, (list(self),)
+
     def __hash__(self):
         return hash(tuple(self))
 

pandas/core/index.py

@@ -2109,7 +2109,7 @@ def __contains__(self, key):
     def __reduce__(self):
         """Necessary for making this object picklable"""
         object_state = list(np.ndarray.__reduce__(self))
-        subclass_state = (self.levels, self.labels, self.sortorder, self.names)
+        subclass_state = (list(self.levels), list(self.labels), self.sortorder, list(self.names))
         object_state[2] = (object_state[2], subclass_state)
         return tuple(object_state)
 

pandas/io/api.py

@@ -10,4 +10,4 @@
 from pandas.io.html import read_html
 from pandas.io.sql import read_sql
 from pandas.io.stata import read_stata
-from pandas.io.pickle import read_pickle
+from pandas.io.pickle import read_pickle, to_pickle

pandas/io/pickle.py

@@ -39,8 +39,7 @@ def try_read(path, encoding=None):
         # the param
         try:
             with open(path,'rb') as fh:
-                with open(path,'rb') as fh:
-                    return pc.load(fh, encoding=encoding, compat=False)
+                return pc.load(fh, encoding=encoding, compat=False)
         except:
             with open(path,'rb') as fh:
                 return pc.load(fh, encoding=encoding, compat=True)

pandas/io/tests/data/legacy_pickle/0.11.0/0.11.0_x86_64_linux_3.3.0.pickle

@@ -77,16 +77,23 @@ def create_data():
 
     index = dict(int = Index(np.arange(10)),
                   date = date_range('20130101',periods=10))
-    mi = dict(reg = MultiIndex.from_tuples(list(zip([['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux'],
+    mi = dict(reg2 = MultiIndex.from_tuples(tuple(zip(*[['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux'],
                                                       ['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two']])),
                                                  names=['first', 'second']))
     series = dict(float = Series(data['A']),
                   int = Series(data['B']),
                   mixed = Series(data['E']),
-                  ts = TimeSeries(np.arange(10).astype(np.int64),index=date_range('20130101',periods=10)))
+                  ts = TimeSeries(np.arange(10).astype(np.int64),index=date_range('20130101',periods=10)),
+                  mi = Series(np.arange(5).astype(np.float64),index=MultiIndex.from_tuples(tuple(zip(*[[1,1,2,2,2],
+                                                                                                    [3,4,3,4,5]])),
+                                                                                           names=['one','two'])))
     frame = dict(float = DataFrame(dict(A = series['float'], B = series['float'] + 1)),
                  int = DataFrame(dict(A = series['int']  , B = series['int']   + 1)),
-                 mixed = DataFrame(dict([ (k,data[k]) for k in ['A','B','C','D']])))
+                 mixed = DataFrame(dict([ (k,data[k]) for k in ['A','B','C','D']])),
+                 mi = DataFrame(dict(A = np.arange(5).astype(np.float64), B = np.arange(5).astype(np.int64)),
+                                index=MultiIndex.from_tuples(tuple(zip(*[['bar','bar','baz','baz','baz'],
+                                                                       ['one','two','one','two','three']])),
+                                                             names=['first','second'])))
     panel = dict(float = Panel(dict(ItemA = frame['float'], ItemB = frame['float']+1)))
 
 

pandas/io/tests/data/legacy_pickle/0.12.0/0.12.0_AMD64_windows_2.7.3.pickle

@@ -15,15 +15,34 @@
 from pandas import Index
 from pandas.sparse.tests import test_sparse
 from pandas import compat
+from pandas.compat import u
 from pandas.util.misc import is_little_endian
 import pandas
 
+def _read_pickle(vf, encoding=None, compat=False):
+    from pandas.compat import pickle_compat as pc
+    with open(vf,'rb') as fh:
+        pc.load(fh, encoding=encoding, compat=compat)
+
 class TestPickle(unittest.TestCase):
     _multiprocess_can_split_ = True
 
     def setUp(self):
         from pandas.io.tests.generate_legacy_pickles import create_data
         self.data = create_data()
+        self.path = u('__%s__.pickle' % tm.rands(10))
+
+    def compare_element(self, typ, result, expected):
+        if isinstance(expected,Index):
+            self.assert_(expected.equals(result))
+            return
+
+        if typ.startswith('sp_'):
+            comparator = getattr(test_sparse,"assert_%s_equal" % typ)
+            comparator(result,expected,exact_indices=False)
+        else:
+            comparator = getattr(tm,"assert_%s_equal" % typ)
+            comparator(result,expected)
 
     def compare(self, vf):
 
@@ -36,19 +55,12 @@ def compare(self, vf):
 
         for typ, dv in data.items():
             for dt, result in dv.items():
-
-                expected = self.data[typ][dt]
-
-                if isinstance(expected,Index):
-                    self.assert_(expected.equals(result))
+                try:
+                    expected = self.data[typ][dt]
+                except (KeyError):
                     continue
 
-                if typ.startswith('sp_'):
-                    comparator = getattr(test_sparse,"assert_%s_equal" % typ)
-                    comparator(result,expected,exact_indices=False)
-                else:
-                    comparator = getattr(tm,"assert_%s_equal" % typ)
-                    comparator(result,expected)
+                self.compare_element(typ, result, expected)
 
     def read_pickles(self, version):
         if not is_little_endian():
@@ -68,8 +80,18 @@ def test_read_pickles_0_11_0(self):
     def test_read_pickles_0_12_0(self):
         self.read_pickles('0.12.0')
 
-    def test_read_pickles_0_13_0(self):
-        self.read_pickles('0.13.0')
+    def test_round_trip_current(self):
+
+        for typ, dv in self.data.items():
+
+            for dt, expected in dv.items():
+
+                with tm.ensure_clean(self.path) as path:
+
+                    pd.to_pickle(expected,path)
+
+                    result = pd.read_pickle(path)
+                    self.compare_element(typ, result, expected)
 
 if __name__ == '__main__':
     import nose

pandas/io/tests/data/legacy_pickle/0.12.0/0.12.0_i686_linux_2.7.3.pickle

@@ -527,7 +527,6 @@ def pxd(name):
                                   'tests/data/legacy_pickle/0.10.1/*.pickle',
                                   'tests/data/legacy_pickle/0.11.0/*.pickle',
                                   'tests/data/legacy_pickle/0.12.0/*.pickle',
-                                  'tests/data/legacy_pickle/0.13.0/*.pickle',
                                   'tests/data/*.csv',
                                   'tests/data/*.dta',
                                   'tests/data/*.txt',