added __getstate__ method (#228)

* added __getstate__ method * fixed lint issues * replaced try except with getattr Co-authored-by: ragrawal <[email protected]>
scikit-learn-contrib · Jan 6, 2021 · 46611ca · 46611ca
1 parent 2296362
commit 46611ca
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/sklearn_pandas/dataframe_mapper.py b/sklearn_pandas/dataframe_mapper.py
@@ -166,6 +166,19 @@ def __setstate__(self, state):
         self.built_default = state.get('built_default', self.default)
         self.transformed_names_ = state.get('transformed_names_', [])
 
+    def __getstate__(self):
+        state = super().__getstate__()
+        state['features'] = self.features
+        state['sparse'] = self.sparse
+        state['default'] = self.default
+        state['df_out'] = self.df_out
+        state['input_df'] = self.input_df
+        state['drop_cols'] = self.drop_cols
+        state['build_features'] = getattr(self, 'built_features', None)
+        state['built_default'] = self.built_default
+        state['transformed_names_'] = self.transformed_names_
+        return state
+
     def _get_col_subset(self, X, cols, input_df=False):
         """
         Get a subset of columns from the given table X.