diff --git a/patsy/contrasts.py b/patsy/contrasts.py index 4173f2e..17c7f5b 100644 --- a/patsy/contrasts.py +++ b/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: diff --git a/patsy/eval.py b/patsy/eval.py index c853988..22ec4dd 100644 --- a/patsy/eval.py +++ b/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): @@ -576,10 +560,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 +576,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 +702,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 +761,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 diff --git a/patsy/test_pickling.py b/patsy/test_pickling.py index fb34667..b666d6d 100644 --- a/patsy/test_pickling.py +++ b/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, diff --git a/pickle_testcases/0.5/contrast_matrix.pickle b/pickle_testcases/0.5/contrast_matrix.pickle index 0d63acd..44ac6f1 100644 Binary files a/pickle_testcases/0.5/contrast_matrix.pickle and b/pickle_testcases/0.5/contrast_matrix.pickle differ diff --git a/pickle_testcases/0.5/evalenv_pickle.pickle b/pickle_testcases/0.5/evalenv_pickle.pickle deleted file mode 100644 index 8e62cdc..0000000 Binary files a/pickle_testcases/0.5/evalenv_pickle.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_simple.pickle b/pickle_testcases/0.5/evalenv_simple.pickle deleted file mode 100644 index 7b8258a..0000000 Binary files a/pickle_testcases/0.5/evalenv_simple.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_bs.pickle b/pickle_testcases/0.5/evalenv_transform_bs.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_bs.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_categorical.pickle b/pickle_testcases/0.5/evalenv_transform_categorical.pickle deleted file mode 100644 index 33d403c..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_categorical.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_cc.pickle b/pickle_testcases/0.5/evalenv_transform_cc.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_cc.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_center.pickle b/pickle_testcases/0.5/evalenv_transform_center.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_center.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_cr.pickle b/pickle_testcases/0.5/evalenv_transform_cr.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_cr.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_scale.pickle b/pickle_testcases/0.5/evalenv_transform_scale.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_scale.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_standardize.pickle b/pickle_testcases/0.5/evalenv_transform_standardize.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_standardize.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/evalenv_transform_te.pickle b/pickle_testcases/0.5/evalenv_transform_te.pickle deleted file mode 100644 index 1d48cbd..0000000 Binary files a/pickle_testcases/0.5/evalenv_transform_te.pickle and /dev/null differ diff --git a/pickle_testcases/0.5/subterm_info.pickle b/pickle_testcases/0.5/subterm_info.pickle index 3b301ff..a09808c 100644 Binary files a/pickle_testcases/0.5/subterm_info.pickle and b/pickle_testcases/0.5/subterm_info.pickle differ diff --git a/pickle_testcases/0.5/transform_center.pickle b/pickle_testcases/0.5/transform_center.pickle index f98f8d2..369a585 100644 Binary files a/pickle_testcases/0.5/transform_center.pickle and b/pickle_testcases/0.5/transform_center.pickle differ diff --git a/pickle_testcases/0.5/transform_standardize_norescale.pickle b/pickle_testcases/0.5/transform_standardize_norescale.pickle index 7218243..d3b0fad 100644 Binary files a/pickle_testcases/0.5/transform_standardize_norescale.pickle and b/pickle_testcases/0.5/transform_standardize_norescale.pickle differ diff --git a/pickle_testcases/0.5/transform_standardize_rescale.pickle b/pickle_testcases/0.5/transform_standardize_rescale.pickle index 7218243..d3b0fad 100644 Binary files a/pickle_testcases/0.5/transform_standardize_rescale.pickle and b/pickle_testcases/0.5/transform_standardize_rescale.pickle differ diff --git a/pickle_testcases/0.5/varlookupdict_simple.pickle b/pickle_testcases/0.5/varlookupdict_simple.pickle deleted file mode 100644 index 853cfcd..0000000 Binary files a/pickle_testcases/0.5/varlookupdict_simple.pickle and /dev/null differ