Adding __getstate__ and __setstate__ on relevant factors. Adding __eq__ as appropriate to make assert tests work.

Author: thequackdaddy

patsy/categorical.py

@@ -46,20 +46,37 @@
                         pandas_Categorical_categories,
                         pandas_Categorical_codes,
                         safe_issubdtype,
-                        no_pickling, assert_no_pickling)
+                        no_pickling, assert_no_pickling, check_pickle_version)
+from patsy.state import StatefulTransform
 
 if have_pandas:
     import pandas
 
 # Objects of this type will always be treated as categorical, with the
 # specified levels and contrast (if given).
+
 class _CategoricalBox(object):
     def __init__(self, data, contrast, levels):
         self.data = data
         self.contrast = contrast
         self.levels = levels
 
-    __getstate__ = no_pickling
+    def __getstate__(self):
+        data = getattr(self, 'data')
+        contrast = getattr(self, 'contrast')
+        levels = getattr(self, 'levels')
+        return (0, data, contrast, levels)
+
+    def __setstate__(self, pickle):
+        version, data, contrast, levels = pickle
+        check_pickle_version(version, 0, name=self.__class__.__name__)
+        self.data = data
+        self.contrast = contrast
+        self.levels = levels
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
 
 def C(data, contrast=None, levels=None):
     """
@@ -120,7 +137,20 @@ def test_C():
     assert c4.contrast == "NEW CONTRAST"
     assert c4.levels == "LEVELS"
 
-    assert_no_pickling(c4)
+    # assert_no_pickling(c4)
+
+
+def test_C_pickle():
+    from six.moves import cPickle as pickle
+    c1 = C("asdf")
+    assert c1 == pickle.loads(pickle.dumps(c1))
+    c2 = C("DATA", "CONTRAST", "LEVELS")
+    assert c2 == pickle.loads(pickle.dumps(c2))
+    c3 = C(c2, levels="NEW LEVELS")
+    assert c3 == pickle.loads(pickle.dumps(c3))
+    c4 = C(c2, "NEW CONTRAST")
+    assert c4 == pickle.loads(pickle.dumps(c4))
+
 
 def guess_categorical(data):
     if safe_is_pandas_categorical(data):
@@ -217,7 +247,7 @@ def sniff(self, data):
         # would be too. Otherwise we need to keep looking.
         return self._level_set == set([True, False])
 
-    __getstate__ = no_pickling
+    # __getstate__ = no_pickling
 
 def test_CategoricalSniffer():
     from patsy.missing import NAAction

patsy/constraint.py

@@ -65,7 +65,10 @@ def _repr_pretty_(self, p, cycle):
         return repr_pretty_impl(p, self,
                                 [self.variable_names, self.coefs, self.constants])
 
-    __getstate__ = no_pickling
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    # __getstate__ = no_pickling
 
     @classmethod
     def combine(cls, constraints):
@@ -118,7 +121,7 @@ def test_LinearConstraint():
     assert_raises(ValueError, LinearConstraint, ["a", "b"],
                   np.zeros((0, 2)))
 
-    assert_no_pickling(lc)
+    # assert_no_pickling(lc)
 
 def test_LinearConstraint_combine():
     comb = LinearConstraint.combine([LinearConstraint(["a", "b"], [1, 0]),

patsy/contrasts.py

@@ -16,7 +16,7 @@
 from patsy import PatsyError
 from patsy.util import (repr_pretty_delegate, repr_pretty_impl,
                         safe_issubdtype,
-                        no_pickling, assert_no_pickling)
+                        no_pickling, assert_no_pickling, check_pickle_version)
 
 class ContrastMatrix(object):
     """A simple container for a matrix used for coding categorical factors.
@@ -47,7 +47,23 @@ def __init__(self, matrix, column_suffixes):
     def _repr_pretty_(self, p, cycle):
         repr_pretty_impl(p, self, [self.matrix, self.column_suffixes])
 
-    __getstate__ = no_pickling
+
+    def __getstate__(self):
+        return (0, self.matrix, self.column_suffixes)
+
+    def __setstate__(self, pickle):
+        version, matrix, column_suffixes = pickle
+        check_pickle_version(version, 0, name=self.__class__.__name__)
+        self.matrix = matrix
+        self.column_suffixes = column_suffixes
+
+    def __eq__(self, other):
+        if self.column_suffixes != other.column_suffixes:
+            return False
+        if not np.array_equal(self.matrix, other.matrix):
+            return False
+        return True
+
 
 def test_ContrastMatrix():
     cm = ContrastMatrix([[1, 0], [0, 1]], ["a", "b"])
@@ -59,7 +75,7 @@ def test_ContrastMatrix():
     from nose.tools import assert_raises
     assert_raises(PatsyError, ContrastMatrix, [[1], [0]], ["a", "b"])
 
-    assert_no_pickling(cm)
+    # assert_no_pickling(cm)
 
 # This always produces an object of the type that Python calls 'str' (whether
 # that be a Python 2 string-of-bytes or a Python 3 string-of-unicode). It does

patsy/desc.py

@@ -14,7 +14,7 @@
 from patsy.eval import EvalEnvironment, EvalFactor
 from patsy.util import uniqueify_list
 from patsy.util import repr_pretty_delegate, repr_pretty_impl
-from patsy.util import no_pickling, assert_no_pickling
+from patsy.util import no_pickling, assert_no_pickling, check_pickle_version
 
 # These are made available in the patsy.* namespace
 __all__ = ["Term", "ModelDesc", "INTERCEPT"]
@@ -65,17 +65,33 @@ def name(self):
         else:
             return "Intercept"
 
-    __getstate__ = no_pickling
+    def __getstate__(self):
+        return (0, self.factors)
+
+    def __setstate__(self, pickle):
+        version, factors = pickle
+        check_pickle_version(version, 0, name=self.__class__.__name__)
+        self.factors = factors
+
+    # __getstate__ = no_pickling
 
 INTERCEPT = Term([])
 
+
 class _MockFactor(object):
     def __init__(self, name):
         self._name = name
 
     def name(self):
         return self._name
 
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __hash__(self):
+        return hash((_MockFactor, str(self._name)))
+
+
 def test_Term():
     assert Term([1, 2, 1]).factors == (1, 2)
     assert Term([1, 2]) == Term([2, 1])
@@ -86,7 +102,11 @@ def test_Term():
     assert Term([f2, f1]).name() == "b:a"
     assert Term([]).name() == "Intercept"
 
-    assert_no_pickling(Term([]))
+    # assert_no_pickling(Term([]))
+
+    from six.moves import cPickle as pickle
+    t = Term([f1, f2])
+    assert t == pickle.loads(pickle.dumps(t, pickle.HIGHEST_PROTOCOL))
 
 class ModelDesc(object):
     """A simple container representing the termlists parsed from a formula.
@@ -166,7 +186,10 @@ def from_formula(cls, tree_or_string):
         assert isinstance(value, cls)
         return value
 
-    __getstate__ = no_pickling
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    # __getstate__ = no_pickling
 
 def test_ModelDesc():
     f1 = _MockFactor("a")
@@ -177,7 +200,9 @@ def test_ModelDesc():
     print(m.describe())
     assert m.describe() == "1 + a ~ 0 + a + a:b"
 
-    assert_no_pickling(m)
+    # assert_no_pickling(m)
+    from six.moves import cPickle as pickle
+    assert m == pickle.loads(pickle.dumps(m, pickle.HIGHEST_PROTOCOL))
 
     assert ModelDesc([], []).describe() == "~ 0"
     assert ModelDesc([INTERCEPT], []).describe() == "1 ~ 0"
@@ -209,7 +234,7 @@ def _pretty_repr_(self, p, cycle): # pragma: no cover
                                 [self.intercept, self.intercept_origin,
                                  self.intercept_removed, self.terms])
 
-    __getstate__ = no_pickling
+    # __getstate__ = no_pickling
 
 def _maybe_add_intercept(doit, terms):
     if doit:

patsy/design_info.py

@@ -32,10 +32,11 @@
 from patsy.compat import OrderedDict
 from patsy.util import (repr_pretty_delegate, repr_pretty_impl,
                         safe_issubdtype,
-                        no_pickling, assert_no_pickling)
+                        no_pickling, assert_no_pickling, check_pickle_version)
 from patsy.constraint import linear_constraint
 from patsy.contrasts import ContrastMatrix
 from patsy.desc import ModelDesc, Term
+from patsy import __version__
 
 class FactorInfo(object):
     """A FactorInfo object is a simple class that provides some metadata about
@@ -120,7 +121,17 @@ def __repr__(self):
             kwlist.append(("categories", self.categories))
         repr_pretty_impl(p, self, [], kwlist)
 
-    __getstate__ = no_pickling
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __hash__(self):
+        if not self.categories:
+            categories = 'NoCategories'
+        else:
+            categories = frozenset(self.categories)
+        return hash((FactorInfo, str(self.factor), str(self.type),
+                    str(self.state), str(self.num_columns), categories))
+
 
 def test_FactorInfo():
     fi1 = FactorInfo("asdf", "numerical", {"a": 1}, num_columns=10)
@@ -234,7 +245,10 @@ def _repr_pretty_(self, p, cycle):
                           ("contrast_matrices", self.contrast_matrices),
                           ("num_columns", self.num_columns)])
 
-    __getstate__ = no_pickling
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    # __getstate__ = no_pickling
 
 def test_SubtermInfo():
     cm = ContrastMatrix(np.ones((2, 2)), ["[1]", "[2]"])
@@ -691,16 +705,40 @@ def from_array(cls, array_like, default_column_prefix="column"):
                             for i in columns]
         return DesignInfo(column_names)
 
-    __getstate__ = no_pickling
+    def __getstate__(self):
+        return (0, self.column_name_indexes, self.factor_infos,
+                self.term_codings, self.term_slices, self.term_name_slices)
+
+    def __setstate__(self, pickle):
+        (version, column_name_indexes, factor_infos, term_codings,
+         term_slices, term_name_slices) = pickle
+        check_pickle_version(version, 0, self.__class__.__name__)
+        self.column_name_indexes = column_name_indexes
+        self.factor_infos = factor_infos
+        self.term_codings = term_codings
+        self.term_slices = term_slices
+        self.term_name_slices = term_name_slices
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+
+class _MockFactor(object):
+    def __init__(self, name):
+        self._name = name
+
+    def name(self):
+        return self._name
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __hash__(self):
+        return hash((_MockFactor, str(self._name)))
+
 
 def test_DesignInfo():
     from nose.tools import assert_raises
-    class _MockFactor(object):
-        def __init__(self, name):
-            self._name = name
-
-        def name(self):
-            return self._name
     f_x = _MockFactor("x")
     f_y = _MockFactor("y")
     t_x = Term([f_x])
@@ -734,8 +772,9 @@ def name(self):
 
     # smoke test
     repr(di)
+    from six.moves import cPickle as pickle
 
-    assert_no_pickling(di)
+    assert di == pickle.loads(pickle.dumps(di, pickle.HIGHEST_PROTOCOL))
 
     # One without term objects
     di = DesignInfo(["a1", "a2", "a3", "b"])
@@ -756,6 +795,8 @@ def name(self):
     assert di.slice("a3") == slice(2, 3)
     assert di.slice("b") == slice(3, 4)
 
+    assert di == pickle.loads(pickle.dumps(di, pickle.HIGHEST_PROTOCOL))
+
     # Check intercept handling in describe()
     assert DesignInfo(["Intercept", "a", "b"]).describe() == "1 + a + b"