Fix weakref support in native classes (#11)

* Fix weakref support in native classes

* Update util.py

Author: BobTheBuidler

mypyc/codegen/emitclass.py

@@ -331,7 +331,7 @@ def emit_line() -> None:
         if emitter.capi_version < (3, 12):
             fields["tp_dictoffset"] = base_size
             fields["tp_weaklistoffset"] = weak_offset
-    elif cl.supports_weakref:
+    elif cl.supports_weakref and emitter.capi_version < (3, 12):
         # __weakref__ lives right after the struct
         # TODO: It should get a member in the struct instead of doing this nonsense.
         emitter.emit_lines(
@@ -340,10 +340,11 @@ def emit_line() -> None:
             "{0}",
             "};",
         )
-        if emitter.capi_version < (3, 12):
-            # versions >= 3.12 set Py_TPFLAGS_MANAGED_WEAKREF flag instead
-            # https://docs.python.org/3.12/extending/newtypes.html#weak-reference-support
-            fields["tp_weaklistoffset"] = base_size
+        fields["tp_members"] = members_name
+        fields["tp_basicsize"] = f"{base_size} + sizeof(PyObject *)"
+        # versions >= 3.12 set Py_TPFLAGS_MANAGED_WEAKREF flag instead
+        # https://docs.python.org/3.12/extending/newtypes.html#weak-reference-support
+        fields["tp_weaklistoffset"] = base_size
     else:
         fields["tp_basicsize"] = base_size
 
@@ -902,6 +903,12 @@ def generate_traverse_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -
             f"*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({struct_name})))",
             object_rprimitive,
         )
+    elif cl.supports_weakref and emitter.capi_version < (3, 12):
+        struct_name = cl.struct_name(emitter.names)
+        # __weakref__ lives right after the struct
+        emitter.emit_gc_visit(
+            f"*((PyObject **)((char *)self + sizeof({struct_name})))", object_rprimitive
+        )
     emitter.emit_line("return 0;")
     emitter.emit_line("}")
 
@@ -925,6 +932,12 @@ def generate_clear_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -> N
             f"*((PyObject **)((char *)self + sizeof(PyObject *) + sizeof({struct_name})))",
             object_rprimitive,
         )
+    elif cl.supports_weakref and emitter.capi_version < (3, 12):
+        struct_name = cl.struct_name(emitter.names)
+        # __weakref__ lives right after the struct
+        emitter.emit_gc_clear(
+            f"*((PyObject **)((char *)self + sizeof({struct_name})))", object_rprimitive
+        )
     emitter.emit_line("return 0;")
     emitter.emit_line("}")
 
@@ -940,12 +953,7 @@ def generate_dealloc_for_class(
     emitter.emit_line(f"{dealloc_func_name}({cl.struct_name(emitter.names)} *self)")
     emitter.emit_line("{")
     if cl.supports_weakref:
-        if emitter.capi_version < (3, 12):
-            emitter.emit_line("if (self->weakreflist != NULL) {")
-            emitter.emit_line("PyObject_ClearWeakRefs((PyObject *) self);")
-            emitter.emit_line("}")
-        else:
-            emitter.emit_line("PyObject_ClearWeakRefs((PyObject *) self);")
+        emitter.emit_line("PyObject_ClearWeakRefs((PyObject *) self);")
     if has_tp_finalize:
         emitter.emit_line("PyObject *type, *value, *traceback;")
         emitter.emit_line("PyErr_Fetch(&type, &value, &traceback);")

mypyc/ir/class_ir.py

@@ -447,7 +447,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         ir.is_final_class = data["is_final_class"]
         ir.inherits_python = data["inherits_python"]
         ir.has_dict = data["has_dict"]
-        ir.supports_weakref = data["supports_weakref"]
+        ir.supports_weakref = data.get("supports_weakref", False)
         ir.allow_interpreted_subclasses = data["allow_interpreted_subclasses"]
         ir.needs_getseters = data["needs_getseters"]
         ir._serializable = data["_serializable"]

mypyc/irbuild/util.py

@@ -33,21 +33,32 @@
 from mypyc.errors import Errors
 
 MYPYC_ATTRS: Final[frozenset[MypycAttr]] = frozenset(
-    ["native_class", "allow_interpreted_subclasses", "serializable", "free_list_len"]
+    [
+        "native_class",
+        "allow_interpreted_subclasses",
+        "serializable",
+        "supports_weakref",
+        "free_list_len",
+    ]
 )
 
 DATACLASS_DECORATORS: Final = frozenset(["dataclasses.dataclass", "attr.s", "attr.attrs"])
 
 
 MypycAttr = Literal[
-    "native_class", "allow_interpreted_subclasses", "serializable", "free_list_len"
+    "native_class",
+    "allow_interpreted_subclasses",
+    "serializable",
+    "supports_weakref",
+    "free_list_len",
 ]
 
 
 class MypycAttrs(TypedDict):
     native_class: NotRequired[bool]
     allow_interpreted_subclasses: NotRequired[bool]
     serializable: NotRequired[bool]
+    supports_weakref: NotRequired[bool]
     free_list_len: NotRequired[int]
 
 

mypyc/test-data/run-weakref.test

@@ -15,6 +15,11 @@ class Object:
 
 _callback_called_cache = {"ref": False, "proxy": False}
 
+@mypyc_attr(supports_weakref=True)
+class NativeObject:
+    def some_meth(self) -> int:
+        return 1
+
 def test_weakref_ref() -> None:
     obj: Optional[Object] = Object()
     r = ref(obj)
@@ -50,3 +55,10 @@ def test_weakref_proxy_with_callback() -> None:
     with assertRaises(ReferenceError):
         p.some_meth()
     assert _callback_called_cache["proxy"] is True
+
+def test_weakref_native_ref() -> None:
+    obj: Optional[NativeObject] = NativeObject()
+    r = ref(obj)
+    assert r() is obj
+    obj = None
+    assert r() is None