BUG: Support WeakSets and ABCMeta instances.

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"""

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()