Abstract Base Classes [Support ABC and Enums - Part 2]

dill/_dill.py

@@ -1700,6 +1700,27 @@ def _get_typedict_type(cls, clsdict, attrs, postproc_list):
         # clsdict.pop('__prepare__', None)
     return clsdict, attrs
 
+def _get_typedict_abc(obj, _dict, attrs, postproc_list):
+    if hasattr(abc, '_get_dump'):
+        (registry, _, _, _) = abc._get_dump(obj)
+        register = obj.register
+        postproc_list.extend((register, (reg(),)) for reg in registry)
+    elif hasattr(obj, '_abc_registry'):
+        registry = obj._abc_registry
+        register = obj.register
+        postproc_list.extend((register, (reg,)) for reg in registry)
+    else:
+        raise PicklingError("Cannot find registry of ABC %s", obj)
+
+    if '_abc_registry' in _dict:
+        del _dict['_abc_registry']
+        del _dict['_abc_cache']
+        del _dict['_abc_negative_cache']
+        # del _dict['_abc_negative_cache_version']
+    else:
+        del _dict['_abc_impl']
+    return _dict, attrs
+
 @register(TypeType)
 def save_type(pickler, obj, postproc_list=None):
     if obj in _typemap:
@@ -1768,6 +1789,11 @@ def save_type(pickler, obj, postproc_list=None):
             for name in slots:
                 _dict.pop(name, None)
 
+            if isinstance(obj, abc.ABCMeta):
+                logger.trace(pickler, "ABC: %s", obj)
+                _dict, attrs = _get_typedict_abc(obj, _dict, attrs, postproc_list)
+                logger.trace(pickler, "# ABC")
+
             qualname = getattr(obj, '__qualname__', None)
             if attrs is not None:
                 for k, v in attrs.items():

dill/tests/test_abc.py

@@ -0,0 +1,163 @@
+#!/usr/bin/env python
+"""
+test dill's ability to pickle abstract base class objects
+"""
+import dill
+import abc
+from abc import ABC
+import warnings
+
+from types import FunctionType
+
+dill.settings['recurse'] = True
+
+class OneTwoThree(ABC):
+    @abc.abstractmethod
+    def foo(self):
+        """A method"""
+        pass
+
+    @property
+    @abc.abstractmethod
+    def bar(self):
+        """Property getter"""
+        pass
+
+    @bar.setter
+    @abc.abstractmethod
+    def bar(self, value):
+        """Property setter"""
+        pass
+
+    @classmethod
+    @abc.abstractmethod
+    def cfoo(cls):
+        """Class method"""
+        pass
+
+    @staticmethod
+    @abc.abstractmethod
+    def sfoo():
+        """Static method"""
+        pass
+
+class EasyAsAbc(OneTwoThree):
+    def __init__(self):
+        self._bar = None
+
+    def foo(self):
+        return "Instance Method FOO"
+
+    @property
+    def bar(self):
+        return self._bar
+
+    @bar.setter
+    def bar(self, value):
+        self._bar = value
+
+    @classmethod
+    def cfoo(cls):
+        return "Class Method CFOO"
+
+    @staticmethod
+    def sfoo():
+        return "Static Method SFOO"
+
+def test_abc_non_local():
+    assert dill.copy(OneTwoThree) is not OneTwoThree
+    assert dill.copy(EasyAsAbc) is not EasyAsAbc
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", dill.PicklingWarning)
+        assert dill.copy(OneTwoThree, byref=True) is OneTwoThree
+        assert dill.copy(EasyAsAbc, byref=True) is EasyAsAbc
+
+    instance = EasyAsAbc()
+    # Set a property that StockPickle can't preserve
+    instance.bar = lambda x: x**2
+    depickled = dill.copy(instance)
+    assert type(depickled) is not type(instance)
+    assert type(depickled.bar) is FunctionType
+    assert depickled.bar(3) == 9
+    assert depickled.sfoo() == "Static Method SFOO"
+    assert depickled.cfoo() == "Class Method CFOO"
+    assert depickled.foo() == "Instance Method FOO"
+
+def test_abc_local():
+    """
+    Test using locally scoped ABC class
+    """
+    class LocalABC(ABC):
+        @abc.abstractmethod
+        def foo(self):
+            pass
+
+        def baz(self):
+            return repr(self)
+
+    labc = dill.copy(LocalABC)
+    assert labc is not LocalABC
+    assert type(labc) is type(LocalABC)
+    # TODO should work like it does for non local classes
+    # <class '__main__.LocalABC'>
+    # <class '__main__.test_abc_local.<locals>.LocalABC'>
+
+    class Real(labc):
+        def foo(self):
+            return "True!"
+
+        def baz(self):
+            return "My " + super(Real, self).baz()
+
+    real = Real()
+    assert real.foo() == "True!"
+
+    try:
+        labc()
+    except TypeError as e:
+        # Expected error
+        pass
+    else:
+        print('Failed to raise type error')
+        assert False
+
+    labc2, pik = dill.copy((labc, Real()))
+    assert 'Real' == type(pik).__name__
+    assert '.Real' in type(pik).__qualname__
+    assert type(pik) is not Real
+    assert labc2 is not LocalABC
+    assert labc2 is not labc
+    assert isinstance(pik, labc2)
+    assert not isinstance(pik, labc)
+    assert not isinstance(pik, LocalABC)
+    assert pik.baz() == "My " + repr(pik)
+
+def test_meta_local_no_cache():
+    """
+    Test calling metaclass and cache registration
+    """
+    LocalMetaABC = abc.ABCMeta('LocalMetaABC', (), {})
+
+    class ClassyClass:
+        pass
+
+    class KlassyClass:
+      pass
+
+    LocalMetaABC.register(ClassyClass)
+
+    assert not issubclass(KlassyClass, LocalMetaABC)
+    assert issubclass(ClassyClass, LocalMetaABC)
+
+    res = dill.dumps((LocalMetaABC, ClassyClass, KlassyClass))
+
+    lmabc, cc, kc = dill.loads(res)
+    assert type(lmabc) == type(LocalMetaABC)
+    assert not issubclass(kc, lmabc)
+    assert issubclass(cc, lmabc)
+
+if __name__ == '__main__':
+    test_abc_non_local()
+    test_abc_local()
+    test_meta_local_no_cache()