From 3f2964de857be51d521852a3b0ae2c80df02f669 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Sat, 1 Jul 2017 16:36:02 -0400 Subject: [PATCH 1/3] BUG: Support WeakSets and ABCMeta instances. Fixes a bug where pickling instances of instances of ABCMeta (you read that correctly) would fail because ABCMeta uses several (previously unpicklable) WeakSets as caches. We fix the issue by adding support for pickling WeakSets via `save_reduce`. --- cloudpickle/cloudpickle.py | 7 +++++ tests/cloudpickle_test.py | 62 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 9f8b3c567..4792889f9 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 sys.version_info >= (2, 7): + 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..cebfd62a3 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 = sys.version_info >= (2, 7) + + def pickle_depickle(obj): """Helper function to test whether object pickled with cloudpickle can be depickled with pickle @@ -592,6 +598,62 @@ 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!') + + # It should still be invalid to construct an instance of the abstract + # class without implementing its methods. + with self.assertRaises(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() From 57f4abc1fda855d70d7d709d94befbefbe813f5d Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Sat, 1 Jul 2017 17:12:12 -0400 Subject: [PATCH 2/3] BUG: assertRaises isn't a contextmanager on 2.6. --- tests/cloudpickle_test.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index cebfd62a3..7e2d4f782 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -619,10 +619,8 @@ def foo(self): self.assertEqual(depickled_class().foo(), 'it works!') self.assertEqual(depickled_instance.foo(), 'it works!') - # It should still be invalid to construct an instance of the abstract - # class without implementing its methods. - with self.assertRaises(TypeError): - depickled_base() + # assertRaises doesn't return a contextmanager in python 2.6 :(. + self.failUnlessRaises(TypeError, depickled_base) class DepickledBaseSubclass(depickled_base): def foo(self): From 8e90a278c1b1a8eca17309e19494933ac27bac10 Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Sat, 1 Jul 2017 17:50:04 -0400 Subject: [PATCH 3/3] MAINT: Use more direct check for WeakSet. --- cloudpickle/cloudpickle.py | 2 +- tests/cloudpickle_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py index 4792889f9..908725e36 100644 --- a/cloudpickle/cloudpickle.py +++ b/cloudpickle/cloudpickle.py @@ -859,7 +859,7 @@ def save_not_implemented(self, obj): dispatch[type(NotImplemented)] = save_not_implemented # WeakSet was added in 2.7. - if sys.version_info >= (2, 7): + if hasattr(weakref, 'WeakSet'): def save_weakset(self, obj): self.save_reduce(weakref.WeakSet, (list(obj),)) diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py index 7e2d4f782..22f656a77 100644 --- a/tests/cloudpickle_test.py +++ b/tests/cloudpickle_test.py @@ -46,7 +46,7 @@ from .testutils import subprocess_pickle_echo -HAVE_WEAKSET = sys.version_info >= (2, 7) +HAVE_WEAKSET = hasattr(weakref, 'WeakSet') def pickle_depickle(obj):