fix: use custom_op

Author: BobTheBuidler

mypyc/codegen/emit.py

@@ -48,6 +48,7 @@
     is_str_rprimitive,
     is_tuple_rprimitive,
     is_uint8_rprimitive,
+    is_weakref_rprimitive,
     object_rprimitive,
     optional_value_type,
 )
@@ -704,7 +705,7 @@ def emit_cast(
             self.emit_lines(f"    {dest} = {src};", "else {")
             self.emit_cast_error_handler(error, src, dest, typ, raise_exception)
             self.emit_line("}")
-        elif is_object_rprimitive(typ):
+        elif is_object_rprimitive(typ) or is_weakref_rprimitive(typ):
             if declare_dest:
                 self.emit_line(f"PyObject *{dest};")
             self.emit_arg_check(src, dest, typ, "", optional)

mypyc/ir/rtypes.py

@@ -637,6 +637,10 @@ def is_sequence_rprimitive(rtype: RType) -> bool:
     )
 
 
+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/specialize.py

@@ -31,6 +31,7 @@
     RefExpr,
     StrExpr,
     TupleExpr,
+    Var,
 )
 from mypy.types import AnyType, TypeOfAny
 from mypyc.ir.ops import (
@@ -103,6 +104,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
@@ -140,6 +142,15 @@ 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 (
+        isinstance(callee, NameExpr)
+        and isinstance(callee.node, Var)
+        # NOTE: why is this not a weakref rprimitive?
+        # TODO: fix to weakref rprimitive so _apply_specialization can use the custom_op
+        and str(callee.node.type).startswith("weakref.ReferenceType")
+        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/misc_ops.c

@@ -1098,3 +1098,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,6 +1,6 @@
 from mypyc.ir.ops import ERR_MAGIC
 from mypyc.ir.rtypes import object_rprimitive, pointer_rprimitive, weakref_rprimitive
-from mypyc.primitives.registry import ERR_NEG_INT, function_op, method_op
+from mypyc.primitives.registry import custom_op, function_op
 
 # Weakref operations
 
@@ -21,10 +21,10 @@
     error_kind=ERR_MAGIC,
 )
 
-deref_op = method_op(
-    name="__call__",
+# TODO: generate specialized versions of this that return the properr rtype
+weakref_deref_op = custom_op(
     arg_types=[weakref_rprimitive],
     return_type=object_rprimitive,
-    c_function_name="PyWeakref_GetRef",
-    error_kind=ERR_NEG_INT,
+    c_function_name="CPyWeakref_GetRef",
+    error_kind=ERR_MAGIC,
 )

mypyc/test-data/irbuild-weakref.test

@@ -2,62 +2,83 @@
 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 :: object
     r0 :: weakref.ReferenceType
-    r1 :: object
+    ref :: object
+    r1 :: weakref.ReferenceType
+    r2 :: object
 L0:
     r0 = PyWeakref_NewRef(x, 0)
-    r1 = PyWeakref_GetRef(r0)
-    return r1
+    ref = r0
+    r1 = cast(weakref.ReferenceType, ref)
+    r2 = CPyWeakref_GetRef(r1)
+    return r2
 
 [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 :: object
     r0 :: weakref.ReferenceType
-    r1 :: object
+    ref :: object
+    r1 :: weakref.ReferenceType
+    r2 :: object
 L0:
     r0 = PyWeakref_NewRef(x, cb)
-    r1 = PyWeakref_GetRef(r0)
-    return r1
+    ref = r0
+    r1 = cast(weakref.ReferenceType, ref)
+    r2 = CPyWeakref_GetRef(r1)
+    return r2
 
 [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 :: object
     r0 :: weakref.ReferenceType
-    r1 :: object
+    r :: object
+    r1 :: weakref.ReferenceType
+    r2 :: object
 L0:
     r0 = PyWeakref_NewRef(x, 0)
-    r1 = PyWeakref_GetRef(r0)
-    return r1
+    r = r0
+    r1 = cast(weakref.ReferenceType, r)
+    r2 = CPyWeakref_GetRef(r1)
+    return r2
 
 [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 :: object
     r0 :: weakref.ReferenceType
-    r1 :: object
+    r :: object
+    r1 :: weakref.ReferenceType
+    r2 :: object
 L0:
     r0 = PyWeakref_NewRef(x, cb)
-    r1 = PyWeakref_GetRef(r0)
-    return r1
+    r = r0
+    r1 = cast(weakref.ReferenceType, r)
+    r2 = CPyWeakref_GetRef(r1)
+    return r2
+