Getting closer on pickling

* Now run subset and clean the eval_env after
removing bare funcalls
* Had some straggling tuples
* Altered test so that stateful_transforms won't
be saved in eval_env
* Removed the clean method on EvalEnvironment

patsy/contrasts.py

@@ -49,13 +49,13 @@ def _repr_pretty_(self, p, cycle):
 
 
     def __getstate__(self):
-        return (0, self.matrix, self.column_suffixes)
+        return {'version': 0, 'matrix': self.matrix,
+                'column_suffixes': 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
+        check_pickle_version(pickle['version'], 0, name=self.__class__.__name__)
+        self.matrix = pickle['matrix']
+        self.column_suffixes = pickle['column_suffixes']
 
     def __eq__(self, other):
         if self.column_suffixes != other.column_suffixes:

patsy/eval.py

@@ -25,6 +25,7 @@
 from patsy.tokens import (pretty_untokenize, normalize_token_spacing,
                              python_tokenize)
 from patsy.compat import call_and_wrap_exc
+from nose.tools import assert_raises
 
 def _all_future_flags():
     flags = 0
@@ -71,13 +72,8 @@ def get(self, key, default=None):
     def __repr__(self):
         return "%s(%r)" % (self.__class__.__name__, self._dicts)
 
-    def __getstate__(self):
-        return (0, self._dicts)
+    __getstate__ = no_pickling
 
-    def __setstate__(self, pickle):
-        version, dicts = pickle
-        check_pickle_version(version, 0, name=self.__class__.__name__)
-        self._dicts = dicts
 
 def test_VarLookupDict():
     d1 = {"a": 1}
@@ -262,27 +258,15 @@ def __hash__(self):
                      tuple(self._namespace_ids())))
 
     def __getstate__(self):
-        self.clean()
+        # self.clean()
         namespaces = self._namespaces
         namespaces = _replace_un_pickleable(namespaces)
-        return (0, namespaces, self.flags)
+        return {'version': 0, 'namespaces': namespaces, 'flags': self.flags}
 
     def __setstate__(self, pickle):
-        version, namespaces, flags = pickle
-        check_pickle_version(version, 0, self.__class__.__name__)
-        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
+        check_pickle_version(pickle['version'], 0, self.__class__.__name__)
+        self.flags = pickle['flags']
+        self._namespaces = _return_un_pickleable(pickle['namespaces'])
 
 
 class ObjectHolder(object):
@@ -504,21 +488,6 @@ def test_EvalEnvironment_eq():
     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)
 # This is purely to make the existence of patsy.builtins visible to systems
@@ -576,10 +545,6 @@ def memorize_passes_needed(self, state, eval_env):
 
         eval_env = eval_env.with_outer_namespace(_builtins_dict)
         env_namespace = eval_env.namespace
-        subset_names = [name for name in ast_names(self.code)
-                        if name in env_namespace]
-        eval_env = eval_env.subset(subset_names)
-        state["eval_env"] = eval_env
 
         # example code: == "2 * center(x)"
         i = [0]
@@ -596,6 +561,12 @@ def new_name_maker(token):
         # example eval_code: == "2 * _patsy_stobj0__center__.transform(x)"
         eval_code = replace_bare_funcalls(self.code, new_name_maker)
         state["eval_code"] = eval_code
+
+        subset_names = [name for name in ast_names(eval_code)
+                        if name in env_namespace]
+        eval_env = eval_env.subset(subset_names)
+        state["eval_env"] = eval_env
+
         # paranoia: verify that none of our new names appeared anywhere in the
         # original code
         if has_bare_variable_reference(state["transforms"], self.code):
@@ -716,7 +687,10 @@ def test_EvalFactor_memorize_passes_needed():
     print(state)
     assert passes == 2
     for name in ["foo", "bar", "quux"]:
-        assert state["eval_env"].namespace[name] is locals()[name]
+        # name should be locally defined, but since its a stateful_transform,
+        # its unnecessary to keep it in eval_env
+        assert name in locals()
+        assert_raises(KeyError, state["eval_env"].namespace.__getitem__, name)
     for name in ["w", "x", "y", "z", "e", "state"]:
         assert name not in state["eval_env"].namespace
     assert state["transforms"] == {"_patsy_stobj0__foo__": "FOO-OBJ",
@@ -772,7 +746,9 @@ def test_EvalFactor_end_to_end():
     print(passes)
     print(state)
     assert passes == 2
-    assert state["eval_env"].namespace["foo"] is foo
+    # We don't want to save the stateful transforms in the eval_env, actually.
+    # Just
+    assert_raises(KeyError, state["eval_env"].namespace.__getitem__, 'foo')
     for name in ["x", "y", "e", "state"]:
         assert name not in state["eval_env"].namespace
     import numpy as np

patsy/test_pickling.py

@@ -42,19 +42,6 @@ def _unwrap_stateful_function(function, *args, **kwargs):
 
 pickling_testcases = {
     "evalfactor_simple": EvalFactor("a+b"),
-    "varlookupdict_simple": VarLookupDict([{"a": 1}, {"a": 2, "b": 3}]),
-    "evalenv_simple": EvalEnvironment([{"a": 1}]),
-    "evalenv_transform_center": EvalEnvironment([{'center': center}]),
-    "evalenv_transform_scale": EvalEnvironment([{'scale': scale}]),
-    "evalenv_transform_standardize": EvalEnvironment([{
-            'standardize': standardize
-    }]),
-    "evalenv_transform_categorical": EvalEnvironment([{'C': C}]),
-    "evalenv_transform_bs": EvalEnvironment([{'cs': bs}]),
-    "evalenv_transform_te": EvalEnvironment([{'te': te}]),
-    "evalenv_transform_cr": EvalEnvironment([{'cs': cr}]),
-    "evalenv_transform_cc": EvalEnvironment([{'cc': cc}]),
-    "evalenv_pickle": EvalEnvironment([{'np': np}]),
     "term": Term([1, 2, 1]),
     "contrast_matrix": ContrastMatrix([[1, 0], [0, 1]], ["a", "b"]),
     "subterm_info": si,