From d2f916c539967150ae21240b7c2922023cd215ac Mon Sep 17 00:00:00 2001
From: Anirudh Vegesana <anirudh.vegesana@gmail.com>
Date: Sun, 8 May 2022 03:42:09 -0700
Subject: [PATCH] Fix class cell in `__new__` for metaclasses (#468)

* Fix class cell in `__new__` for metaclasses

* Python 3 super
---
 dill/_dill.py          |  8 +++++++-
 tests/test_classdef.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/dill/_dill.py b/dill/_dill.py
index 40c01fc5..1c3caaed 100644
--- a/dill/_dill.py
+++ b/dill/_dill.py
@@ -1589,7 +1589,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
diff --git a/tests/test_classdef.py b/tests/test_classdef.py
index ca6ab6ad..e23bef0a 100644
--- a/tests/test_classdef.py
+++ b/tests/test_classdef.py
@@ -216,6 +216,41 @@ def test_slots():
     assert dill.pickles(Y.y)
     assert dill.copy(y).y == value
 
+def test_metaclass():
+    if dill._dill.PY3:
+        class metaclass_with_new(type):
+            def __new__(mcls, name, bases, ns, **kwds):
+                cls = super().__new__(mcls, name, bases, ns, **kwds)
+                assert mcls is not None
+                assert cls.method(mcls)
+                return cls
+            def method(cls, mcls):
+                return isinstance(cls, mcls)
+
+        l = locals()
+        exec("""class subclass_with_new(metaclass=metaclass_with_new):
+            def __new__(cls):
+                self = super().__new__(cls)
+                return self""", None, l)
+        subclass_with_new = l['subclass_with_new']
+    else:
+        class metaclass_with_new(type):
+            def __new__(mcls, name, bases, ns, **kwds):
+                cls = super(mcls, metaclass_with_new).__new__(mcls, name, bases, ns, **kwds)
+                assert mcls is not None
+                assert cls.method(mcls)
+                return cls
+            def method(cls, mcls):
+                return isinstance(cls, mcls)
+
+        class subclass_with_new:
+            __metaclass__ = metaclass_with_new
+            def __new__(cls):
+                self = super(subclass_with_new, cls).__new__(cls)
+                return self
+
+    assert dill.copy(subclass_with_new())
+
 
 if __name__ == '__main__':
     test_class_instances()
@@ -227,3 +262,4 @@ def test_slots():
     test_array_subclass()
     test_method_decorator()
     test_slots()
+    test_metaclass()