feat: new mypyc primitives for weakref.proxy (#19217)

BobTheBuidler · web-flow · commit 8f2371a565eb · 2025-09-05T16:15:36.000+01:00
This PR adds 2 new weakref primitives for weakref.proxy (1 and 2 arg)
diff --git a/mypyc/primitives/weakref_ops.py b/mypyc/primitives/weakref_ops.py
@@ -20,3 +20,21 @@
     c_function_name="PyWeakref_NewRef",
     error_kind=ERR_MAGIC,
 )
+
+new_proxy_op = function_op(
+    name="_weakref.proxy",
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="PyWeakref_NewProxy",
+    extra_int_constants=[(0, pointer_rprimitive)],
+    error_kind=ERR_MAGIC,
+)
+
+new_proxy_with_callback_op = function_op(
+    name="_weakref.proxy",
+    arg_types=[object_rprimitive, object_rprimitive],
+    # steals=[True, False],
+    return_type=object_rprimitive,
+    c_function_name="PyWeakref_NewProxy",
+    error_kind=ERR_MAGIC,
+)
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -343,6 +343,7 @@ class RuntimeError(Exception): pass
 class UnicodeEncodeError(RuntimeError): pass
 class UnicodeDecodeError(RuntimeError): pass
 class NotImplementedError(RuntimeError): pass
+class ReferenceError(Exception): pass
 
 class StopIteration(Exception):
     value: Any
diff --git a/mypyc/test-data/irbuild-weakref.test b/mypyc/test-data/irbuild-weakref.test
@@ -49,3 +49,55 @@ def f(x, cb):
 L0:
     r0 = PyWeakref_NewRef(x, cb)
     return r0
+
+[case testWeakrefProxy]
+import weakref
+from typing import Any, Callable
+def f(x: object) -> object:
+    return weakref.proxy(x)
+
+[out]
+def f(x):
+    x, r0 :: object
+L0:
+    r0 = PyWeakref_NewProxy(x, 0)
+    return r0
+
+[case testWeakrefProxyCallback]
+import weakref
+from typing import Any, Callable
+def f(x: object, cb: Callable[[object], Any]) -> object:
+    return weakref.proxy(x, cb)
+
+[out]
+def f(x, cb):
+    x, cb, r0 :: object
+L0:
+    r0 = PyWeakref_NewProxy(x, cb)
+    return r0
+
+[case testFromWeakrefProxy]
+from typing import Any, Callable
+from weakref import proxy
+def f(x: object) -> object:
+    return proxy(x)
+
+[out]
+def f(x):
+    x, r0 :: object
+L0:
+    r0 = PyWeakref_NewProxy(x, 0)
+    return r0
+
+[case testFromWeakrefProxyCallback]
+from typing import Any, Callable
+from weakref import proxy
+def f(x: object, cb: Callable[[object], Any]) -> object:
+    return proxy(x, cb)
+
+[out]
+def f(x, cb):
+    x, cb, r0 :: object
+L0:
+    r0 = PyWeakref_NewProxy(x, cb)
+    return r0
diff --git a/mypyc/test-data/run-weakref.test b/mypyc/test-data/run-weakref.test
@@ -1,30 +1,52 @@
 # Test cases for weakrefs (compile and run)
 
 [case testWeakrefRef]
-from weakref import ref
+# mypy: disable-error-code="union-attr"
+from weakref import proxy, ref
 from mypy_extensions import mypyc_attr
+from testutil import assertRaises
+from typing import Optional
 
 @mypyc_attr(native_class=False)
 class Object:
     """some random weakreffable object"""
-    pass
+    def some_meth(self) -> int:
+        return 1
 
-def test_weakref_ref():
-    obj = Object()
+_callback_called_cache = {"ref": False, "proxy": False}
+
+def test_weakref_ref() -> None:
+    obj: Optional[Object] = Object()
     r = ref(obj)
     assert r() is obj
     obj = None
     assert r() is None, r()
 
-def test_weakref_ref_with_callback():
-    obj = Object()
-    r = ref(obj, lambda x: x)
+def test_weakref_ref_with_callback() -> None:
+    obj: Optional[Object] = Object()
+    r = ref(obj, lambda x: _callback_called_cache.__setitem__("ref", True))
     assert r() is obj
     obj = None
     assert r() is None, r()
+    assert _callback_called_cache["ref"] is True
 
-[file driver.py]
-from native import test_weakref_ref, test_weakref_ref_with_callback
+def test_weakref_proxy() -> None:
+    obj: Optional[Object] = Object()
+    p = proxy(obj)
+    assert obj.some_meth() == 1
+    assert p.some_meth() == 1
+    obj.some_meth()
+    obj = None
+    with assertRaises(ReferenceError):
+        p.some_meth()
 
-test_weakref_ref()
-test_weakref_ref_with_callback()
+def test_weakref_proxy_with_callback() -> None:
+    obj: Optional[Object] = Object()
+    p = proxy(obj, lambda x: _callback_called_cache.__setitem__("proxy", True))
+    assert obj.some_meth() == 1
+    assert p.some_meth() == 1
+    obj.some_meth()
+    obj = None
+    with assertRaises(ReferenceError):
+        p.some_meth()
+    assert _callback_called_cache["proxy"] is True
diff --git a/test-data/unit/lib-stub/_weakref.pyi b/test-data/unit/lib-stub/_weakref.pyi
@@ -0,0 +1,11 @@
+from typing import Any, Callable, TypeVar, overload
+from weakref import CallableProxyType, ProxyType
+
+_C = TypeVar("_C", bound=Callable[..., Any])
+_T = TypeVar("_T")
+
+# Return CallableProxyType if object is callable, ProxyType otherwise
+@overload
+def proxy(object: _C, callback: Callable[[CallableProxyType[_C]], Any] | None = None, /) -> CallableProxyType[_C]: ...
+@overload
+def proxy(object: _T, callback: Callable[[ProxyType[_T]], Any] | None = None, /) -> ProxyType[_T]: ...
diff --git a/test-data/unit/lib-stub/weakref.pyi b/test-data/unit/lib-stub/weakref.pyi
@@ -1,11 +1,23 @@
+from _weakref import proxy
 from collections.abc import Callable
-from typing import Any, Generic, TypeVar
+from typing import Any, ClassVar, Generic, TypeVar, final
 from typing_extensions import Self
 
+_C = TypeVar("_C", bound=Callable[..., Any])
 _T = TypeVar("_T")
 
 class ReferenceType(Generic[_T]):  # "weakref"
     __callback__: Callable[[Self], Any]
     def __new__(cls, o: _T, callback: Callable[[Self], Any] | None = ..., /) -> Self: ...
+    def __call__(self) -> _T | None: ...
 
 ref = ReferenceType
+
+@final
+class CallableProxyType(Generic[_C]):  # "weakcallableproxy"
+    def __eq__(self, value: object, /) -> bool: ...
+    def __getattr__(self, attr: str) -> Any: ...
+    __call__: _C
+    __hash__: ClassVar[None]  # type: ignore[assignment]
+
+__all__ = ["proxy"]
