diff options
| -rw-r--r-- | lib/sqlalchemy/orm/attributes.py | 6 | ||||
| -rw-r--r-- | lib/sqlalchemy/orm/interfaces.py | 3 | ||||
| -rw-r--r-- | test/orm/pickled.py | 19 |
3 files changed, 26 insertions, 2 deletions
diff --git a/lib/sqlalchemy/orm/attributes.py b/lib/sqlalchemy/orm/attributes.py index 2b2760208..5c242aa7e 100644 --- a/lib/sqlalchemy/orm/attributes.py +++ b/lib/sqlalchemy/orm/attributes.py @@ -1269,6 +1269,9 @@ class ClassManager(dict): def setup_instance(self, instance, state=None): setattr(instance, self.STATE_ATTR, state or self.instance_state_factory(instance, self)) + + def teardown_instance(self, instance): + delattr(instance, self.STATE_ATTR) def _new_state_if_none(self, instance): """Install a default InstanceState if none is present. @@ -1381,6 +1384,9 @@ class _ClassInstrumentationAdapter(ClassManager): state.dict = self._adapted.get_instance_dict(self.class_, instance) return state + def teardown_instance(self, instance): + self._adapted.remove_state(self.class_, instance) + def state_of(self, instance): if hasattr(self._adapted, 'state_of'): return self._adapted.state_of(self.class_, instance) diff --git a/lib/sqlalchemy/orm/interfaces.py b/lib/sqlalchemy/orm/interfaces.py index 25664f258..b210e577f 100644 --- a/lib/sqlalchemy/orm/interfaces.py +++ b/lib/sqlalchemy/orm/interfaces.py @@ -889,5 +889,8 @@ class InstrumentationManager(object): def install_state(self, class_, instance, state): setattr(instance, '_default_state', state) + def remove_state(self, class_, instance): + delattr(instance, '_default_state', state) + def state_getter(self, class_): return lambda instance: getattr(instance, '_default_state') diff --git a/test/orm/pickled.py b/test/orm/pickled.py index a4eadf5a9..d14acdc50 100644 --- a/test/orm/pickled.py +++ b/test/orm/pickled.py @@ -2,7 +2,7 @@ import testenv; testenv.configure_for_tests() import pickle from testlib import sa, testing from testlib.sa import Table, Column, Integer, String, ForeignKey -from testlib.sa.orm import mapper, relation, create_session +from testlib.sa.orm import mapper, relation, create_session, attributes from orm import _base, _fixtures @@ -141,6 +141,21 @@ class PolymorphicDeferredTest(_base.MappedTest): assert 'email_address' not in eu2.__dict__ self.assertEquals(eu2.email_address, 'foo@bar.com') - +class CustomSetupTeardowntest(_fixtures.FixtureTest): + @testing.resolve_artifact_names + def test_rebuild_state(self): + """not much of a 'test', but illustrate how to + remove instance-level state before pickling. + + """ + mapper(User, users) + + u1 = User() + attributes.manager_of_class(User).teardown_instance(u1) + assert not u1.__dict__ + u2 = pickle.loads(pickle.dumps(u1)) + attributes.manager_of_class(User).setup_instance(u2) + assert attributes.instance_state(u2) + if __name__ == '__main__': testenv.main() |