[mypyc] feat: new primitive for weakref.ref.__call__

mypyc/ir/rtypes.py

@@ -512,6 +512,12 @@ def __hash__(self) -> int:
 # Python range object.
 range_rprimitive: Final = RPrimitive("builtins.range", is_unboxed=False, is_refcounted=True)
 
+# Python weak reference object
+weakref_rprimitive: Final = RPrimitive(
+    "weakref.ReferenceType", is_unboxed=False, is_refcounted=True
+)
+
+
 KNOWN_NATIVE_TYPES: Final = {
     name: RPrimitive(name, is_unboxed=False, is_refcounted=True)
     for name in ["native_internal.Buffer"]
@@ -653,6 +659,10 @@ def is_immutable_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
     )
 
 
+def is_weakref_rprimitive(rtype: RType) -> TypeGuard[RPrimitive]:
+    return isinstance(rtype, RPrimitive) and rtype.name == "weakref.ReferenceType"
+
+
 class TupleNameVisitor(RTypeVisitor[str]):
     """Produce a tuple name based on the concrete representations of types."""
 

mypyc/irbuild/mapper.py

@@ -47,6 +47,7 @@
     str_rprimitive,
     tuple_rprimitive,
     uint8_rprimitive,
+    weakref_rprimitive,
 )
 
 
@@ -103,6 +104,8 @@ def type_to_rtype(self, typ: Type | None) -> RType:
                 return tuple_rprimitive  # Varying-length tuple
             elif typ.type.fullname == "builtins.range":
                 return range_rprimitive
+            elif typ.type.fullname == "weakref.ReferenceType":
+                return weakref_rprimitive
             elif typ.type in self.type_to_ir:
                 inst = RInstance(self.type_to_ir[typ.type])
                 # Treat protocols as Union[protocol, object], so that we can do fast

mypyc/irbuild/specialize.py

@@ -70,6 +70,7 @@
     is_int_rprimitive,
     is_list_rprimitive,
     is_uint8_rprimitive,
+    is_weakref_rprimitive,
     list_rprimitive,
     object_rprimitive,
     set_rprimitive,
@@ -114,6 +115,7 @@
     str_encode_utf8_strict,
 )
 from mypyc.primitives.tuple_ops import isinstance_tuple, new_tuple_set_item_op
+from mypyc.primitives.weakref_ops import weakref_deref_op
 
 # Specializers are attempted before compiling the arguments to the
 # function.  Specializers can return None to indicate that they failed
@@ -151,6 +153,8 @@ def apply_function_specialization(
     builder: IRBuilder, expr: CallExpr, callee: RefExpr
 ) -> Value | None:
     """Invoke the Specializer callback for a function if one has been registered"""
+    if is_weakref_rprimitive(builder.node_type(callee)) and len(expr.args) == 0:
+        return builder.call_c(weakref_deref_op, [builder.accept(expr.callee)], expr.line)
     return _apply_specialization(builder, expr, callee, callee.fullname)
 
 

mypyc/lib-rt/CPy.h

@@ -946,6 +946,7 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, PyOb
 
 PyObject *CPy_GetAIter(PyObject *obj);
 PyObject *CPy_GetANext(PyObject *aiter);
+PyObject *CPyWeakref_GetRef(PyObject *ref);
 void CPy_SetTypeAliasTypeComputeFunction(PyObject *alias, PyObject *compute_value);
 void CPyTrace_LogEvent(const char *location, const char *line, const char *op, const char *details);
 

mypyc/lib-rt/misc_ops.c

@@ -1138,3 +1138,18 @@ void CPy_SetImmortal(PyObject *obj) {
 }
 
 #endif
+
+
+PyObject *CPyWeakref_GetRef(PyObject *ref)
+{
+    PyObject *obj = NULL;
+    int success = PyWeakref_GetRef(ref, &obj);
+    if (success == -1) {
+        return NULL;
+    } else if (obj == NULL) {
+        Py_INCREF(Py_None);
+        return Py_None;
+    } else {
+        return obj;
+    }
+}

mypyc/primitives/weakref_ops.py

@@ -1,13 +1,13 @@
 from mypyc.ir.ops import ERR_MAGIC
-from mypyc.ir.rtypes import object_rprimitive, pointer_rprimitive
-from mypyc.primitives.registry import function_op
+from mypyc.ir.rtypes import object_rprimitive, pointer_rprimitive, weakref_rprimitive
+from mypyc.primitives.registry import custom_op, function_op
 
 # Weakref operations
 
 new_ref_op = function_op(
     name="weakref.ReferenceType",
     arg_types=[object_rprimitive],
-    return_type=object_rprimitive,
+    return_type=weakref_rprimitive,
     c_function_name="PyWeakref_NewRef",
     extra_int_constants=[(0, pointer_rprimitive)],
     error_kind=ERR_MAGIC,
@@ -16,7 +16,7 @@
 new_ref__with_callback_op = function_op(
     name="weakref.ReferenceType",
     arg_types=[object_rprimitive, object_rprimitive],
-    return_type=object_rprimitive,
+    return_type=weakref_rprimitive,
     c_function_name="PyWeakref_NewRef",
     error_kind=ERR_MAGIC,
 )
@@ -38,3 +38,11 @@
     c_function_name="PyWeakref_NewProxy",
     error_kind=ERR_MAGIC,
 )
+
+# TODO: generate specialized versions of this that return the proper rtype
+weakref_deref_op = custom_op(
+    arg_types=[weakref_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPyWeakref_GetRef",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/irbuild-weakref.test

@@ -2,53 +2,73 @@
 import weakref
 from typing import Any, Callable
 def f(x: object) -> object:
-    return weakref.ref(x)
+    ref = weakref.ref(x)
+    return ref()
 
 [out]
 def f(x):
-    x, r0 :: object
+    x :: object
+    r0, ref :: weakref.ReferenceType
+    r1 :: object
 L0:
     r0 = PyWeakref_NewRef(x, 0)
-    return r0
+    ref = r0
+    r1 = CPyWeakref_GetRef(ref)
+    return r1
 
 [case testWeakrefRefCallback]
 import weakref
 from typing import Any, Callable
 def f(x: object, cb: Callable[[object], Any]) -> object:
-    return weakref.ref(x, cb)
+    ref = weakref.ref(x, cb)
+    return ref()
 
 [out]
 def f(x, cb):
-    x, cb, r0 :: object
+    x, cb :: object
+    r0, ref :: weakref.ReferenceType
+    r1 :: object
 L0:
     r0 = PyWeakref_NewRef(x, cb)
-    return r0
+    ref = r0
+    r1 = CPyWeakref_GetRef(ref)
+    return r1
 
 [case testFromWeakrefRef]
 from typing import Any, Callable
 from weakref import ref
 def f(x: object) -> object:
-    return ref(x)
+    r = ref(x)
+    return r()
 
 [out]
 def f(x):
-    x, r0 :: object
+    x :: object
+    r0, r :: weakref.ReferenceType
+    r1 :: object
 L0:
     r0 = PyWeakref_NewRef(x, 0)
-    return r0
+    r = r0
+    r1 = CPyWeakref_GetRef(r)
+    return r1
 
 [case testFromWeakrefRefCallback]
 from typing import Any, Callable
 from weakref import ref
 def f(x: object, cb: Callable[[object], Any]) -> object:
-    return ref(x, cb)
+    r = ref(x, cb)
+    return r()
 
 [out]
 def f(x, cb):
-    x, cb, r0 :: object
+    x, cb :: object
+    r0, r :: weakref.ReferenceType
+    r1 :: object
 L0:
     r0 = PyWeakref_NewRef(x, cb)
-    return r0
+    r = r0
+    r1 = CPyWeakref_GetRef(r)
+    return r1
 
 [case testWeakrefProxy]
 import weakref