Merge pull request #10557 from kjordahl/bug/pickle-subclass-metadata

Pickle subclass metadata

Author: jreback

doc/source/whatsnew/v0.17.0.txt

@@ -51,6 +51,7 @@ Other API Changes
 - Enable writing Excel files in :ref:`memory <_io.excel_writing_buffer>` using StringIO/BytesIO (:issue:`7074`)
 - Enable serialization of lists and dicts to strings in ExcelWriter (:issue:`8188`)
 - Allow passing `kwargs` to the interpolation methods (:issue:`10378`).
+- Serialize metadata properties of subclasses of pandas objects (:issue:`10553`).
 
 .. _whatsnew_0170.deprecations:
 

pandas/core/generic.py

@@ -80,8 +80,8 @@ class NDFrame(PandasObject):
     """
     _internal_names = ['_data', '_cacher', '_item_cache', '_cache',
                        'is_copy', '_subtyp', '_index',
-                       '_default_kind', '_default_fill_value',
-                       '__array_struct__','__array_interface__']
+                       '_default_kind', '_default_fill_value', '_metadata',
+                       '__array_struct__', '__array_interface__']
     _internal_names_set = set(_internal_names)
     _accessors = frozenset([])
     _metadata = []
@@ -760,7 +760,9 @@ def to_dense(self):
     # Picklability
 
     def __getstate__(self):
-        return self._data
+        meta = dict((k, getattr(self, k, None)) for k in self._metadata)
+        return dict(_data=self._data, _typ=self._typ,
+                    _metadata=self._metadata, **meta)
 
     def __setstate__(self, state):
 

pandas/core/series.py

@@ -438,10 +438,6 @@ def imag(self, v):
     __long__ = _coerce_method(int)
     __int__ = _coerce_method(int)
 
-    # we are preserving name here
-    def __getstate__(self):
-        return dict(_data=self._data, name=self.name)
-
     def _unpickle_series_compat(self, state):
         if isinstance(state, dict):
             self._data = state['_data']

pandas/tests/test_frame.py

@@ -45,7 +45,8 @@
                                  assertRaisesRegexp,
                                  assertRaises,
                                  makeCustomDataframe as mkdf,
-                                 ensure_clean)
+                                 ensure_clean,
+                                 SubclassedDataFrame)
 from pandas.core.indexing import IndexingError
 from pandas.core.common import PandasError
 
@@ -14501,16 +14502,8 @@ def test_assign_bad(self):
 
     def test_dataframe_metadata(self):
 
-        class TestDataFrame(DataFrame):
-            _metadata = ['testattr']
-
-            @property
-            def _constructor(self):
-                return TestDataFrame
-
-
-        df = TestDataFrame({'X': [1, 2, 3], 'Y': [1, 2, 3]},
-                           index=['a', 'b', 'c'])
+        df = SubclassedDataFrame({'X': [1, 2, 3], 'Y': [1, 2, 3]},
+                                 index=['a', 'b', 'c'])
         df.testattr = 'XXX'
 
         self.assertEqual(df.testattr, 'XXX')
@@ -14519,6 +14512,11 @@ def _constructor(self):
         self.assertEqual(df.iloc[[0, 1], :].testattr, 'XXX')
         # GH9776
         self.assertEqual(df.iloc[0:1, :].testattr, 'XXX')
+        # GH10553
+        unpickled = self.round_trip_pickle(df)
+        assert_frame_equal(df, unpickled)
+        self.assertEqual(df._metadata, unpickled._metadata)
+        self.assertEqual(df.testattr, unpickled.testattr)
 
     def test_to_panel_expanddim(self):
         # GH 9762

pandas/util/testing.py

@@ -1856,3 +1856,11 @@ def inner(*args, **kwargs):
                 thread.join()
         return inner
     return wrapper
+
+
+class SubclassedDataFrame(DataFrame):
+    _metadata = ['testattr']
+
+    @property
+    def _constructor(self):
+        return SubclassedDataFrame