Skip to content

Commit

Permalink
Clean the EvalEnvironment before pickling to removing patsy's statefu…
Browse files Browse the repository at this point in the history
…l transforms which have different names/qualnames from expected.
  • Loading branch information
thequackdaddy committed Nov 4, 2018
1 parent 0301817 commit 019094d
Show file tree
Hide file tree
Showing 11 changed files with 30 additions and 15 deletions.
31 changes: 30 additions & 1 deletion patsy/eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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):
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
6 changes: 0 additions & 6 deletions patsy/mgcv_cubic_splines.py
Original file line number Diff line number Diff line change
Expand Up @@ -730,8 +730,6 @@ def __setstate__(self, pickle):


cr = stateful_transform(CR)
cr.__qualname__ = 'cr'
cr.__name__ = 'cr'


class CC(CubicRegressionSpline):
Expand Down Expand Up @@ -774,8 +772,6 @@ def __setstate__(self, pickle):


cc = stateful_transform(CC)
cc.__qualname__ = 'cc'
cc.__name__ = 'cc'


def test_crs_errors():
Expand Down Expand Up @@ -978,8 +974,6 @@ def __setstate__(self, pickle):


te = stateful_transform(TE)
te.__qualname__ = 'te'
te.__name__ = 'te'


def test_te_errors():
Expand Down
2 changes: 0 additions & 2 deletions patsy/splines.py
Original file line number Diff line number Diff line change
Expand Up @@ -257,8 +257,6 @@ def __setstate__(self, pickle):


bs = stateful_transform(BS)
bs.__qualname__ = 'bs'
bs.__name__ = 'bs'


def test_bs_compat():
Expand Down
6 changes: 0 additions & 6 deletions patsy/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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'
Binary file modified pickle_testcases/0.5/evalenv_transform_bs.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_cc.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_center.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_cr.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_scale.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_standardize.pickle
Binary file not shown.
Binary file modified pickle_testcases/0.5/evalenv_transform_te.pickle
Binary file not shown.

0 comments on commit 019094d

Please sign in to comment.