Clean the EvalEnvironment before pickling to removing patsy's stateful transforms which have different names/qualnames from expected.

Author: thequackdaddy

patsy/eval.py

@@ -262,6 +262,7 @@ def __hash__(self):
                      tuple(self._namespace_ids())))
 
     def __getstate__(self):
+        self.clean()
         namespaces = self._namespaces
         namespaces = _replace_un_pickleable(namespaces)
         return (0, namespaces, self.flags)
@@ -272,6 +273,17 @@ def __setstate__(self, pickle):
         self.flags = flags
         self._namespaces = _return_un_pickleable(namespaces)
 
+    def clean(self):
+        """The EvalEnvironment doesn't need the stateful transformation
+        functions once the design matrix has been built. This will delete
+        it. Called by __getstate__ to prepare for pickling."""
+        namespaces = []
+        for namespace in self._namespaces:
+            ns = {key: namespace[key] for key in six.iterkeys(namespace) if not
+                  hasattr(namespace[key], '__patsy_stateful_transform__')}
+            namespaces.append(ns)
+        self._namespaces = namespaces
+
 
 class ObjectHolder(object):
     def __init__(self, kind, module, name):
@@ -489,7 +501,23 @@ def test_EvalEnvironment_eq():
     capture_local_env = lambda: EvalEnvironment.capture(0)
     env3 = capture_local_env()
     env4 = capture_local_env()
-    assert env3 != env4 # This fails... 
+    assert env3 != env4
+
+
+def test_EvalEnvironment_clean():
+    from patsy.state import center, standardize
+    from patsy.splines import bs
+
+    env1 = EvalEnvironment([{'center': center}])
+    env2 = EvalEnvironment([{'standardize': standardize}])
+    env3 = EvalEnvironment([{'bs': bs}])
+    env1.clean()
+    env2.clean()
+    env3.clean()
+
+    env1._namespaces == [{}]
+    env2._namespaces == [{}]
+    env3._namespaces == [{}]
 
 _builtins_dict = {}
 six.exec_("from patsy.builtins import *", {}, _builtins_dict)
@@ -650,6 +678,7 @@ def __setstate__(self, pickle):
         self.code = pickle['code']
         self.origin = pickle['origin']
 
+
 def test_EvalFactor_pickle_saves_origin():
     from patsy.util import assert_pickled_equals
     # The pickling tests use object equality before and after pickling

patsy/mgcv_cubic_splines.py

@@ -730,8 +730,6 @@ def __setstate__(self, pickle):
 
 
 cr = stateful_transform(CR)
-cr.__qualname__ = 'cr'
-cr.__name__ = 'cr'
 
 
 class CC(CubicRegressionSpline):
@@ -774,8 +772,6 @@ def __setstate__(self, pickle):
 
 
 cc = stateful_transform(CC)
-cc.__qualname__ = 'cc'
-cc.__name__ = 'cc'
 
 
 def test_crs_errors():
@@ -978,8 +974,6 @@ def __setstate__(self, pickle):
 
 
 te = stateful_transform(TE)
-te.__qualname__ = 'te'
-te.__name__ = 'te'
 
 
 def test_te_errors():

patsy/splines.py

@@ -257,8 +257,6 @@ def __setstate__(self, pickle):
 
 
 bs = stateful_transform(BS)
-bs.__qualname__ = 'bs'
-bs.__name__ = 'bs'
 
 
 def test_bs_compat():

patsy/state.py

@@ -126,8 +126,6 @@ def __setstate__(self, pickle):
 
 
 center = stateful_transform(Center)
-center.__qualname__ = 'center'
-center.__name__ = 'center'
 
 # See:
 #   http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#On-line_algorithm
@@ -196,9 +194,5 @@ def __setstate__(self, pickle):
 
 
 standardize = stateful_transform(Standardize)
-standardize.__qualname__ = 'standardize'
-standardize.__name__ = 'standardize'
 # R compatibility:
 scale = standardize
-scale.__qualname__ = 'scale'
-scale.__name__ = 'scale'