diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 9f8b3c567..908725e36 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -858,6 +858,13 @@ def save_not_implemented(self, obj): dispatch[type(Ellipsis)] = save_ellipsis dispatch[type(NotImplemented)] = save_not_implemented + # WeakSet was added in 2.7. + if hasattr(weakref, 'WeakSet'): + def save_weakset(self, obj): + self.save_reduce(weakref.WeakSet, (list(obj),)) + + dispatch[weakref.WeakSet] = save_weakset + """Special functions for Add-on libraries""" def inject_addons(self): """Plug in system. Register additional pickling functions if modules already loaded""" diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index b8ad025ce..22f656a77 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -1,5 +1,7 @@ from __future__ import division +import abc + import base64 import functools import imp @@ -14,6 +16,7 @@ import sys import textwrap import unittest +import weakref try: from StringIO import StringIO @@ -43,6 +46,9 @@ from .testutils import subprocess_pickle_echo +HAVE_WEAKSET = hasattr(weakref, 'WeakSet') + + def pickle_depickle(obj): """Helper function to test whether object pickled with cloudpickle can be depickled with pickle @@ -592,6 +598,60 @@ def test_logger(self): self.assertEqual(out.strip().decode(), 'INFO:cloudpickle.dummy_test_logger:hello') + def test_abc(self): + + @abc.abstractmethod + def foo(self): + raise NotImplementedError('foo') + + # Invoke the metaclass directly rather than using class syntax for + # python 2/3 compat. + AbstractClass = abc.ABCMeta('AbstractClass', (object,), {'foo': foo}) + + class ConcreteClass(AbstractClass): + def foo(self): + return 'it works!' + + depickled_base = pickle_depickle(AbstractClass) + depickled_class = pickle_depickle(ConcreteClass) + depickled_instance = pickle_depickle(ConcreteClass()) + + self.assertEqual(depickled_class().foo(), 'it works!') + self.assertEqual(depickled_instance.foo(), 'it works!') + + # assertRaises doesn't return a contextmanager in python 2.6 :(. + self.failUnlessRaises(TypeError, depickled_base) + + class DepickledBaseSubclass(depickled_base): + def foo(self): + return 'it works for realz!' + + self.assertEqual(DepickledBaseSubclass().foo(), 'it works for realz!') + + @pytest.mark.skipif(not HAVE_WEAKSET, reason="WeakSet doesn't exist") + def test_weakset_identity_preservation(self): + # Test that weaksets don't lose all their inhabitants if they're + # pickled in a larger data structure that includes other references to + # their inhabitants. + + class SomeClass(object): + def __init__(self, x): + self.x = x + + obj1, obj2, obj3 = SomeClass(1), SomeClass(2), SomeClass(3) + + things = [weakref.WeakSet([obj1, obj2]), obj1, obj2, obj3] + result = pickle_depickle(things) + + weakset, depickled1, depickled2, depickled3 = result + + self.assertEqual(depickled1.x, 1) + self.assertEqual(depickled2.x, 2) + self.assertEqual(depickled3.x, 3) + self.assertEqual(len(weakset), 2) + + self.assertEqual(set(weakset), set([depickled1, depickled2])) + if __name__ == '__main__': unittest.main()