Fix ABC.__new__ bug and add a heinous ABCEnum test case

anivegesana · anivegesana · commit 1b6da3695e98 · 2022-05-06T01:24:00.000-07:00
diff --git a/dill/_dill.py b/dill/_dill.py
@@ -1619,7 +1619,13 @@ def save_cell(pickler, obj):
         log.info("# Ce3")
         return
     if is_dill(pickler, child=True):
-        postproc = next(iter(pickler._postproc.values()), None)
+        if id(f) in pickler._postproc:
+            # Already seen. Add to its postprocessing.
+            postproc = pickler._postproc[id(f)]
+        else:
+            # Haven't seen it. Add to the highest possible object and set its
+            # value as late as possible to prevent cycle.
+            postproc = next(iter(pickler._postproc.values()), None)
         if postproc is not None:
             log.info("Ce2: %s" % obj)
             # _CELL_REF is defined in _shims.py to support older versions of
@@ -1830,7 +1836,7 @@ def _get_typedict_type(cls, clsdict, postproc_list):
     return clsdict
     # return _dict_from_dictproxy(cls.__dict__)
 
-def _get_typedict_abc(obj, _dict, state, postproc_list):
+def _get_typedict_abc(obj, _dict, attrs, postproc_list):
     log.info("ABC: %s" % obj)
     if hasattr(abc, '_get_dump'):
         (registry, _, _, _) = abc._get_dump(obj)
@@ -1851,9 +1857,9 @@ def _get_typedict_abc(obj, _dict, state, postproc_list):
     else:
         del _dict['_abc_impl']
     log.info("# ABC")
-    return _dict, state
+    return _dict, attrs
 
-def _get_typedict_enum(obj, _dict, state, postproc_list):
+def _get_typedict_enum(obj, _dict, attrs, postproc_list):
     log.info("E: %s" % obj)
     metacls = type(obj)
     original_dict = {}
@@ -1866,8 +1872,12 @@ def _get_typedict_enum(obj, _dict, state, postproc_list):
     _dict.pop('_value2member_map_', None)
     _dict.pop('_generate_next_value_', None)
 
+    if attrs is not None:
+        attrs.update(_dict)
+        _dict = attrs
+
     log.info("# E")
-    return original_dict, (None, _dict)
+    return original_dict, _dict
 
 @register(TypeType)
 def save_type(pickler, obj, postproc_list=None):
@@ -1912,7 +1922,6 @@ def save_type(pickler, obj, postproc_list=None):
         log.info("# 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)
@@ -1923,25 +1932,30 @@ def save_type(pickler, obj, postproc_list=None):
             # thanks to Tom Stepleton pointing out pickler._session unneeded
             _t = 'T3'
             _dict = _get_typedict_type(obj, obj.__dict__.copy(), postproc_list) # copy dict proxy to a dict
-            state = None
+            attrs = None
 
             for name in _dict.get("__slots__", []):
                 del _dict[name]
 
-            if PY3 and obj_name != obj.__name__:
-                postproc_list.append((setattr, (obj, '__qualname__', obj_name)))
-
             if isinstance(obj, abc.ABCMeta):
-                _dict, state = _get_typedict_abc(obj, _dict, state, postproc_list)
+                _dict, attrs = _get_typedict_abc(obj, _dict, attrs, postproc_list)
 
             if EnumMeta and isinstance(obj, EnumMeta):
-                _dict, state = _get_typedict_enum(obj, _dict, state, postproc_list)
-
-           #print (_dict)
-           #print ("%s\n%s" % (type(obj), obj.__name__))
-           #print ("%s\n%s" % (obj.__bases__, obj.__dict__))
-
-            if PY3 and type(obj) is not type or hasattr(obj, '__orig_bases__'):
+                _dict, attrs = _get_typedict_enum(obj, _dict, attrs, postproc_list)
+
+            qualname = getattr(obj, '__qualname__', None)
+            if attrs is not None:
+                if qualname is not None:
+                    attrs['__qualname__'] = qualname
+                for k, v in attrs.items():
+                    postproc_list.append((setattr, (obj, k, v)))
+                state = _dict, attrs
+            elif qualname is not None:
+                postproc_list.append((setattr, (obj, '__qualname__', qualname)))
+                state = _dict
+
+            if True: # PY3 and type(obj) is not type or hasattr(obj, '__orig_bases__'):
+                # This case will always work, but might be overkill.
                 from types import new_class
                 _metadict = {
                     'metaclass': type(obj)
@@ -1962,6 +1976,7 @@ def save_type(pickler, obj, postproc_list=None):
                 )), state, obj=obj, postproc_list=postproc_list)
             log.info("# %s" % _t)
         else:
+            obj_name = getattr(obj, '__qualname__', getattr(obj, '__name__', None))
             log.info("T4: %s" % obj)
             if incorrectly_named:
                 warnings.warn('Cannot locate reference to %r.' % (obj,), PicklingWarning)
diff --git a/tests/test_enum.py b/tests/test_enum.py
@@ -0,0 +1,138 @@
+try:
+    import enum
+    from enum import Enum, IntEnum, EnumMeta, Flag, IntFlag, unique, auto
+except:
+    Enum = None
+
+import abc
+
+import dill
+import sys
+
+dill.settings['recurse'] = True
+
+"""
+Test cases copied from https://raw.githubusercontent.com/python/cpython/3.10/Lib/test/test_enum.py
+
+Copyright 1991-1995 by Stichting Mathematisch Centrum, Amsterdam, The Netherlands.
+
+All Rights Reserved
+Permission to use, copy, modify, and distribute this software and its
+documentation for any purpose and without fee is hereby granted, provided that
+the above copyright notice appear in all copies and that both that copyright
+notice and this permission notice appear in supporting documentation, and that
+the names of Stichting Mathematisch Centrum or CWI or Corporation for National
+Research Initiatives or CNRI not be used in advertising or publicity pertaining
+to distribution of the software without specific, written prior permission.
+
+While CWI is the initial source for this software, a modified version is made
+available by the Corporation for National Research Initiatives (CNRI) at the
+Internet address http://www.python.org.
+
+STICHTING MATHEMATISCH CENTRUM AND CNRI DISCLAIM ALL WARRANTIES WITH REGARD TO
+THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS,
+IN NO EVENT SHALL STICHTING MATHEMATISCH CENTRUM OR CNRI BE LIABLE FOR ANY
+SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
+FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
+OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+PERFORMANCE OF THIS SOFTWARE.
+"""
+
+def test_enums():
+
+    class Stooges(Enum):
+        LARRY = 1
+        CURLY = 2
+        MOE = 3
+
+    class IntStooges(int, Enum):
+        LARRY = 1
+        CURLY = 2
+        MOE = 3
+
+    class FloatStooges(float, Enum):
+        LARRY = 1.39
+        CURLY = 2.72
+        MOE = 3.142596
+
+    class FlagStooges(Flag):
+        LARRY = 1
+        CURLY = 2
+        MOE = 3
+
+    # https://stackoverflow.com/a/56135108
+    class ABCEnumMeta(abc.ABCMeta, EnumMeta):
+        def __new__(mcls, *args, **kw):
+            abstract_enum_cls = super().__new__(mcls, *args, **kw)
+            # Only check abstractions if members were defined.
+            if abstract_enum_cls._member_map_:
+                try:  # Handle existence of undefined abstract methods.
+                    absmethods = list(abstract_enum_cls.__abstractmethods__)
+                    if absmethods:
+                        missing = ', '.join(f'{method!r}' for method in absmethods)
+                        plural = 's' if len(absmethods) > 1 else ''
+                        raise TypeError(
+                           f"cannot instantiate abstract class {abstract_enum_cls.__name__!r}"
+                           f" with abstract method{plural} {missing}")
+                except AttributeError:
+                    pass
+            return abstract_enum_cls
+
+    if dill._dill.PY3:
+        l = locals()
+        exec("""class StrEnum(str, abc.ABC, Enum, metaclass=ABCEnumMeta):
+            'accepts only string values'
+            def invisible(self):
+                return "did you see me?" """, None, l)
+        StrEnum = l['StrEnum']
+    else:
+        class StrEnum(str, abc.ABC, Enum):
+            __metaclass__ = ABCEnumMeta
+            'accepts only string values'
+            def invisible(self):
+                return "did you see me?"
+
+    class Name(StrEnum):
+        BDFL = 'Guido van Rossum'
+        FLUFL = 'Barry Warsaw'
+
+    assert 'invisible' in dir(dill.copy(Name).BDFL)
+    assert 'invisible' in dir(dill.copy(Name.BDFL))
+    assert dill.copy(Name.BDFL) is not Name.BDFL
+
+    Question = Enum('Question', 'who what when where why', module=__name__)
+    Answer = Enum('Answer', 'him this then there because')
+    Theory = Enum('Theory', 'rule law supposition', qualname='spanish_inquisition')
+
+    class Fruit(Enum):
+        TOMATO = 1
+        BANANA = 2
+        CHERRY = 3
+
+    assert dill.copy(Fruit).TOMATO.value == 1 and dill.copy(Fruit).TOMATO != 1 \
+        and dill.copy(Fruit).TOMATO is not Fruit.TOMATO
+
+    from datetime import date
+    class Holiday(date, Enum):
+        NEW_YEAR = 2013, 1, 1
+        IDES_OF_MARCH = 2013, 3, 15
+        
+    # TODO: Fix this case
+    # assert hasattr(dill.copy(Holiday), 'NEW_YEAR')
+
+    class SuperEnum(IntEnum):
+        def __new__(cls, value, description=""):
+            obj = int.__new__(cls, value)
+            obj._value_ = value
+            obj.description = description
+            return obj
+
+    class SubEnum(SuperEnum):
+        sample = 5
+
+    if sys.hexversion >= 0x030a0000:
+        assert 'description' in dir(dill.copy(SubEnum.sample))
+        assert 'description' in dir(dill.copy(SubEnum).sample)
+
+if __name__ == '__main__':
+    test_enums()
