[mypyc] Fix __dict__ in CPython 3.11

mypyc/codegen/emitclass.py

@@ -257,21 +257,40 @@ def emit_line() -> None:
 
     # XXX: there is no reason for the __weakref__ stuff to be mixed up with __dict__
     if cl.has_dict:
-        # __dict__ lives right after the struct and __weakref__ lives right after that
-        # TODO: They should get members in the struct instead of doing this nonsense.
-        weak_offset = f"{base_size} + sizeof(PyObject *)"
-        emitter.emit_lines(
-            f"PyMemberDef {members_name}[] = {{",
-            f'{{"__dict__", T_OBJECT_EX, {base_size}, 0, NULL}},',
-            f'{{"__weakref__", T_OBJECT_EX, {weak_offset}, 0, NULL}},',
-            "{0}",
-            "};",
-        )
+        # TODO: Use PY_VERSION_HEX and include patchlevel.h
+        if emitter.capi_version < (3, 11):
+            # __dict__ lives right after the struct and __weakref__ lives right after that
+            # TODO: They should get members in the struct instead of doing this nonsense.
+            weak_offset = f"{base_size} + sizeof(PyObject *)"
+            emitter.emit_lines(
+                f"PyMemberDef {members_name}[] = {{",
+                f'{{"__dict__", T_OBJECT_EX, {base_size}, 0, NULL}},',
+                f'{{"__weakref__", T_OBJECT_EX, {weak_offset}, 0, NULL}},',
+                "{0}",
+                "};",
+            )
+
+            fields["tp_members"] = members_name
+            fields["tp_basicsize"] = f"{base_size} + 2*sizeof(PyObject *)"
+            fields["tp_dictoffset"] = base_size
+            fields["tp_weaklistoffset"] = weak_offset
+
+        else:
+            # In CPython >=3.11, __dict__ slot is automatically set before the type. So
+            # we need to let CPython manage it.
+            # __weakref__ lives right after the struct
+            # TODO: __weakref__ should get members in the struct instead of doing this nonsense.
+            weak_offset = f"{base_size}"
+            emitter.emit_lines(
+                f"PyMemberDef {members_name}[] = {{",
+                f'{{"__weakref__", T_OBJECT_EX, {weak_offset}, 0, NULL}},',
+                "{0}",
+                "};",
+            )
 
-        fields["tp_members"] = members_name
-        fields["tp_basicsize"] = f"{base_size} + 2*sizeof(PyObject *)"
-        fields["tp_dictoffset"] = base_size
-        fields["tp_weaklistoffset"] = weak_offset
+            fields["tp_members"] = members_name
+            fields["tp_basicsize"] = f"{base_size} + sizeof(PyObject *)"
+            fields["tp_weaklistoffset"] = weak_offset
     else:
         fields["tp_basicsize"] = base_size
 
@@ -327,6 +346,9 @@ def emit_line() -> None:
             # This is just a placeholder to please CPython. It will be
             # overriden during setup.
             fields["tp_call"] = "PyVectorcall_Call"
+    # TODO: Use PY_VERSION_HEX and include patchlevel.h
+    if cl.has_dict and emitter.capi_version >= (3, 11):
+        flags.append("Py_TPFLAGS_MANAGED_DICT")
     fields["tp_flags"] = " | ".join(flags)
 
     emitter.emit_line(f"static PyTypeObject {emitter.type_struct_name(cl)}_template_ = {{")
@@ -707,14 +729,35 @@ def generate_traverse_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -
             emitter.emit_gc_visit(f"self->{emitter.attr(attr)}", rtype)
     if cl.has_dict:
         struct_name = cl.struct_name(emitter.names)
-        # __dict__ lives right after the struct and __weakref__ lives right after that
-        emitter.emit_gc_visit(
-            "*((PyObject **)((char *)self + sizeof({})))".format(struct_name), object_rprimitive
-        )
-        emitter.emit_gc_visit(
-            "*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({})))".format(struct_name),
-            object_rprimitive,
-        )
+        # TODO: Use PY_VERSION_HEX and include patchlevel.h
+        if emitter.capi_version < (3, 11):
+            # __dict__ lives right after the struct and __weakref__ lives right after that
+            emitter.emit_gc_visit(
+                "*((PyObject **)((char *)self + sizeof({})))".format(struct_name),
+                object_rprimitive,
+            )
+            emitter.emit_gc_visit(
+                "*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({})))".format(
+                    struct_name
+                ),
+                object_rprimitive,
+            )
+        else:
+            # In >=3.11, __dict__ is managed by CPython and __weakref__ lives right after
+            # the struct
+            emitter.emit_line("{")
+            emitter.emit_line(
+                "int vret = _PyObject_VisitManagedDict((PyObject *)self, visit, arg);"
+            )
+            emitter.emit_line("if (vret)")
+            emitter.emit_line("{")
+            emitter.emit_line("return vret;")
+            emitter.emit_line("}")
+            emitter.emit_line("}")
+            emitter.emit_gc_visit(
+                "*((PyObject **)((char *)self + sizeof({})))".format(struct_name),
+                object_rprimitive,
+            )
     emitter.emit_line("return 0;")
     emitter.emit_line("}")
 
@@ -728,14 +771,27 @@ def generate_clear_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -> N
             emitter.emit_gc_clear(f"self->{emitter.attr(attr)}", rtype)
     if cl.has_dict:
         struct_name = cl.struct_name(emitter.names)
-        # __dict__ lives right after the struct and __weakref__ lives right after that
-        emitter.emit_gc_clear(
-            "*((PyObject **)((char *)self + sizeof({})))".format(struct_name), object_rprimitive
-        )
-        emitter.emit_gc_clear(
-            "*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({})))".format(struct_name),
-            object_rprimitive,
-        )
+        # TODO: Use PY_VERSION_HEX and include patchlevel.h
+        if emitter.capi_version < (3, 11):
+            # __dict__ lives right after the struct and __weakref__ lives right after that
+            emitter.emit_gc_clear(
+                "*((PyObject **)((char *)self + sizeof({})))".format(struct_name),
+                object_rprimitive,
+            )
+            emitter.emit_gc_clear(
+                "*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({})))".format(
+                    struct_name
+                ),
+                object_rprimitive,
+            )
+        else:
+            # In >=3.11, __dict__ is managed by CPython and __weakref__ lives right after
+            # the struct
+            emitter.emit_line("_PyObject_ClearManagedDict((PyObject *)self);")
+            emitter.emit_gc_clear(
+                "*((PyObject **)((char *)self + sizeof({})))".format(struct_name),
+                object_rprimitive,
+            )
     emitter.emit_line("return 0;")
     emitter.emit_line("}")