Python 3 Metaclasses [Support ABC and Enums - Part 1] (#577)

* Python 3 Metaclasses

* Add more tests for Python 3 Metaclasses

Author: anivegesana

dill/_dill.py

@@ -65,15 +65,17 @@
 SliceType = slice
 TypeType = type # 'new-style' classes #XXX: unregistered
 XRangeType = range
-from types import MappingProxyType as DictProxyType
+from types import MappingProxyType as DictProxyType, new_class
 from pickle import DEFAULT_PROTOCOL, HIGHEST_PROTOCOL, PickleError, PicklingError, UnpicklingError
 import __main__ as _main_module
 import marshal
 import gc
 # import zlib
+import abc
 import dataclasses
 from weakref import ReferenceType, ProxyType, CallableProxyType
 from collections import OrderedDict
+from enum import Enum, EnumMeta
 from functools import partial
 from operator import itemgetter, attrgetter
 GENERATOR_FAIL = False
@@ -1669,6 +1671,35 @@ def save_module(pickler, obj):
             logger.trace(pickler, "# M2")
     return
 
+# The following function is based on '_extract_class_dict' from 'cloudpickle'
+# Copyright (c) 2012, Regents of the University of California.
+# Copyright (c) 2009 `PiCloud, Inc. <http://www.picloud.com>`_.
+# License: https://github.com/cloudpipe/cloudpickle/blob/master/LICENSE
+def _get_typedict_type(cls, clsdict, attrs, postproc_list):
+    """Retrieve a copy of the dict of a class without the inherited methods"""
+    if len(cls.__bases__) == 1:
+        inherited_dict = cls.__bases__[0].__dict__
+    else:
+        inherited_dict = {}
+        for base in reversed(cls.__bases__):
+            inherited_dict.update(base.__dict__)
+    to_remove = []
+    for name, value in dict.items(clsdict):
+        try:
+            base_value = inherited_dict[name]
+            if value is base_value:
+                to_remove.append(name)
+        except KeyError:
+            pass
+    for name in to_remove:
+        dict.pop(clsdict, name)
+
+    if issubclass(type(cls), type):
+        clsdict.pop('__dict__', None)
+        clsdict.pop('__weakref__', None)
+        # clsdict.pop('__prepare__', None)
+    return clsdict, attrs
+
 @register(TypeType)
 def save_type(pickler, obj, postproc_list=None):
     if obj in _typemap:
@@ -1680,15 +1711,22 @@ def save_type(pickler, obj, postproc_list=None):
     elif obj.__bases__ == (tuple,) and all([hasattr(obj, attr) for attr in ('_fields','_asdict','_make','_replace')]):
         # special case: namedtuples
         logger.trace(pickler, "T6: %s", obj)
+
+        obj_name = getattr(obj, '__qualname__', getattr(obj, '__name__', None))
+        if obj.__name__ != obj_name:
+            if postproc_list is None:
+                postproc_list = []
+            postproc_list.append((setattr, (obj, '__qualname__', obj_name)))
+
         if not obj._field_defaults:
-            pickler.save_reduce(_create_namedtuple, (obj.__name__, obj._fields, obj.__module__), obj=obj)
+            _save_with_postproc(pickler, (_create_namedtuple, (obj.__name__, obj._fields, obj.__module__)), obj=obj, postproc_list=postproc_list)
         else:
             defaults = [obj._field_defaults[field] for field in obj._fields if field in obj._field_defaults]
-            pickler.save_reduce(_create_namedtuple, (obj.__name__, obj._fields, obj.__module__, defaults), obj=obj)
+            _save_with_postproc(pickler, (_create_namedtuple, (obj.__name__, obj._fields, obj.__module__, defaults)), obj=obj, postproc_list=postproc_list)
         logger.trace(pickler, "# T6")
         return
 
-    # special cases: NoneType, NotImplementedType, EllipsisType
+    # special cases: NoneType, NotImplementedType, EllipsisType, EnumMeta
     elif obj is type(None):
         logger.trace(pickler, "T7: %s", obj)
         #XXX: pickler.save_reduce(type, (None,), obj=obj)
@@ -1702,35 +1740,63 @@ def save_type(pickler, obj, postproc_list=None):
         logger.trace(pickler, "T7: %s", obj)
         pickler.save_reduce(type, (Ellipsis,), obj=obj)
         logger.trace(pickler, "# T7")
+    elif obj is EnumMeta:
+        logger.trace(pickler, "T7: %s", obj)
+        pickler.write(GLOBAL + b'enum\nEnumMeta\n')
+        logger.trace(pickler, "# T7")
 
     else:
-        obj_name = getattr(obj, '__qualname__', getattr(obj, '__name__', None))
         _byref = getattr(pickler, '_byref', None)
         obj_recursive = id(obj) in getattr(pickler, '_postproc', ())
         incorrectly_named = not _locate_function(obj, pickler)
         if not _byref and not obj_recursive and incorrectly_named: # not a function, but the name was held over
+            if postproc_list is None:
+                postproc_list = []
+
             # thanks to Tom Stepleton pointing out pickler._session unneeded
             logger.trace(pickler, "T2: %s", obj)
-            _dict = obj.__dict__.copy() # convert dictproxy to dict
+            _dict, attrs = _get_typedict_type(obj, obj.__dict__.copy(), None, postproc_list) # copy dict proxy to a dict
+
            #print (_dict)
            #print ("%s\n%s" % (type(obj), obj.__name__))
            #print ("%s\n%s" % (obj.__bases__, obj.__dict__))
             slots = _dict.get('__slots__', ())
-            if type(slots) == str: slots = (slots,) # __slots__ accepts a single string
+            if type(slots) == str:
+                # __slots__ accepts a single string
+                slots = (slots,)
+
             for name in slots:
-                del _dict[name]
-            _dict.pop('__dict__', None)
-            _dict.pop('__weakref__', None)
-            _dict.pop('__prepare__', None)
-            if obj_name != obj.__name__:
-                if postproc_list is None:
-                    postproc_list = []
-                postproc_list.append((setattr, (obj, '__qualname__', obj_name)))
-            _save_with_postproc(pickler, (_create_type, (
-                type(obj), obj.__name__, obj.__bases__, _dict
-            )), obj=obj, postproc_list=postproc_list)
+                _dict.pop(name, None)
+
+            qualname = getattr(obj, '__qualname__', None)
+            if attrs is not None:
+                for k, v in attrs.items():
+                    postproc_list.append((setattr, (obj, k, v)))
+                # TODO: Consider using the state argument to save_reduce?
+            if qualname is not None:
+                postproc_list.append((setattr, (obj, '__qualname__', qualname)))
+
+            if not hasattr(obj, '__orig_bases__'):
+                _save_with_postproc(pickler, (_create_type, (
+                    type(obj), obj.__name__, obj.__bases__, _dict
+                )), obj=obj, postproc_list=postproc_list)
+            else:
+                # This case will always work, but might be overkill.
+                _metadict = {
+                    'metaclass': type(obj)
+                }
+
+                if _dict:
+                    _dict_update = PartialType(_setitems, source=_dict)
+                else:
+                    _dict_update = None
+
+                _save_with_postproc(pickler, (new_class, (
+                    obj.__name__, obj.__orig_bases__, _metadict, _dict_update
+                )), obj=obj, postproc_list=postproc_list)
             logger.trace(pickler, "# T2")
         else:
+            obj_name = getattr(obj, '__qualname__', getattr(obj, '__name__', None))
             logger.trace(pickler, "T4: %s", obj)
             if incorrectly_named:
                 warnings.warn(
@@ -1753,14 +1819,17 @@ def save_type(pickler, obj, postproc_list=None):
     return
 
 @register(property)
+@register(abc.abstractproperty)
 def save_property(pickler, obj):
     logger.trace(pickler, "Pr: %s", obj)
-    pickler.save_reduce(property, (obj.fget, obj.fset, obj.fdel, obj.__doc__),
+    pickler.save_reduce(type(obj), (obj.fget, obj.fset, obj.fdel, obj.__doc__),
                         obj=obj)
     logger.trace(pickler, "# Pr")
 
 @register(staticmethod)
 @register(classmethod)
+@register(abc.abstractstaticmethod)
+@register(abc.abstractclassmethod)
 def save_classmethod(pickler, obj):
     logger.trace(pickler, "Cm: %s", obj)
     orig_func = obj.__func__

dill/tests/test_classdef.py

@@ -7,6 +7,7 @@
 #  - https://github.com/uqfoundation/dill/blob/master/LICENSE
 
 import dill
+from enum import EnumMeta
 import sys
 dill.settings['recurse'] = True
 
@@ -54,6 +55,14 @@ def ok(self):
 nc = _newclass2()
 m = _mclass()
 
+if sys.hexversion < 0x03090000:
+    import typing
+    class customIntList(typing.List[int]):
+        pass
+else:
+    class customIntList(list[int]):
+        pass
+
 # test pickles for class instances
 def test_class_instances():
     assert dill.pickles(o)
@@ -89,6 +98,7 @@ def test_specialtypes():
     assert dill.pickles(type(None))
     assert dill.pickles(type(NotImplemented))
     assert dill.pickles(type(Ellipsis))
+    assert dill.pickles(type(EnumMeta))
 
 from collections import namedtuple
 Z = namedtuple("Z", ['a','b'])
@@ -99,19 +109,23 @@ def test_specialtypes():
 Xi = X(0,1)
 Bad = namedtuple("FakeName", ['a','b'])
 Badi = Bad(0,1)
+Defaults = namedtuple('Defaults', ['x', 'y'], defaults=[1])
+Defaultsi = Defaults(2)
 
 # test namedtuple
 def test_namedtuple():
     assert Z is dill.loads(dill.dumps(Z))
     assert Zi == dill.loads(dill.dumps(Zi))
     assert X is dill.loads(dill.dumps(X))
     assert Xi == dill.loads(dill.dumps(Xi))
+    assert Defaults is dill.loads(dill.dumps(Defaults))
+    assert Defaultsi == dill.loads(dill.dumps(Defaultsi))
     assert Bad is not dill.loads(dill.dumps(Bad))
     assert Bad._fields == dill.loads(dill.dumps(Bad))._fields
     assert tuple(Badi) == tuple(dill.loads(dill.dumps(Badi)))
 
     class A:
-        class B(namedtuple("B", ["one", "two"])):
+        class B(namedtuple("C", ["one", "two"])):
             '''docstring'''
         B.__module__ = 'testing'
 
@@ -123,6 +137,15 @@ class B(namedtuple("B", ["one", "two"])):
     assert dill.copy(A.B).__doc__ == 'docstring'
     assert dill.copy(A.B).__module__ == 'testing'
 
+    from typing import NamedTuple
+
+    def A():
+        class B(NamedTuple):
+            x: int
+        return B
+
+    assert type(dill.copy(A()(8))).__qualname__ == type(A()(8)).__qualname__
+
 def test_dtype():
     try:
         import numpy as np
@@ -204,11 +227,20 @@ def __init__(self, y):
 value = 123
 y = Y(value)
 
+class Y2(object):
+  __slots__ = 'y'
+  def __init__(self, y):
+    self.y = y
+
 def test_slots():
     assert dill.pickles(Y)
     assert dill.pickles(y)
     assert dill.pickles(Y.y)
     assert dill.copy(y).y == value
+    assert dill.copy(Y2(value)).y == value
+
+def test_origbases():
+    assert dill.copy(customIntList).__orig_bases__ == customIntList.__orig_bases__
 
 def test_attr():
     import attr
@@ -238,6 +270,11 @@ def __new__(cls):
 
     assert dill.copy(subclass_with_new())
 
+def test_enummeta():
+    from http import HTTPStatus
+    import enum
+    assert dill.copy(HTTPStatus.OK) is HTTPStatus.OK
+    assert dill.copy(enum.EnumMeta) is enum.EnumMeta
 
 if __name__ == '__main__':
     test_class_instances()
@@ -249,4 +286,6 @@ def __new__(cls):
     test_array_subclass()
     test_method_decorator()
     test_slots()
+    test_origbases()
     test_metaclass()
+    test_enummeta()