diff --git a/mypyc/primitives/registry.py b/mypyc/primitives/registry.py index 5e7ecb70f55d..07546663d08e 100644 --- a/mypyc/primitives/registry.py +++ b/mypyc/primitives/registry.py @@ -371,4 +371,5 @@ def load_address_op(name: str, type: RType, src: str) -> LoadAddressDescription: import mypyc.primitives.list_ops import mypyc.primitives.misc_ops import mypyc.primitives.str_ops -import mypyc.primitives.tuple_ops # noqa: F401 +import mypyc.primitives.tuple_ops +import mypyc.primitives.weakref_ops # noqa: F401 diff --git a/mypyc/primitives/weakref_ops.py b/mypyc/primitives/weakref_ops.py new file mode 100644 index 000000000000..024f9fedd17f --- /dev/null +++ b/mypyc/primitives/weakref_ops.py @@ -0,0 +1,21 @@ +from mypyc.ir.ops import ERR_MAGIC +from mypyc.ir.rtypes import object_rprimitive +from mypyc.primitives.registry import ERR_NEG_INT, function_op + +# Weakref operations + +new_ref_op = function_op( + name="weakref.ReferenceType", + arg_types=[object_rprimitive, object_rprimitive], + return_type=object_rprimitive, + c_function_name="PyWeakref_NewRef", + error_kind=ERR_MAGIC, +) + +deref_op = function_op( + name="weakref.ReferenceType.__call__", + arg_types=[object_rprimitive], + return_type=object_rprimitive, + c_function_name="PyWeakref_GetRef", + error_kind=ERR_NEG_INT, +) diff --git a/mypyc/test-data/irbuild-weakref.test b/mypyc/test-data/irbuild-weakref.test new file mode 100644 index 000000000000..1b8c894821e5 --- /dev/null +++ b/mypyc/test-data/irbuild-weakref.test @@ -0,0 +1,27 @@ +[case testWeakrefRefCallback] +import weakref +from typing import Any, Callable +def f(x: object, cb: Callable[[object], Any]) -> object: + return weakref.ref(x, cb)() + +[out] +def f(x, cb): + x, cb, r0, r1 :: object +L0: + r0 = PyWeakref_NewRef(x, cb) + r1 = PyWeakref_GetRef(r0) + 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)() + +[out] +def f(x, cb): + x, cb, r0, r1 :: object +L0: + r0 = PyWeakref_NewRef(x, cb) + r1 = PyWeakref_GetRef(r0) + return r1 diff --git a/mypyc/test-data/run-weakref.test b/mypyc/test-data/run-weakref.test new file mode 100644 index 000000000000..de2d5dd5bad2 --- /dev/null +++ b/mypyc/test-data/run-weakref.test @@ -0,0 +1,19 @@ +# Test cases for weakrefs (compile and run) + +[case testWeakrefRef] +import asyncio +import pytest # type: ignore [import-not-found] +from weakref import ref +from mypy_extensions import mypyc_attr + +@mypyc_attr(native_class=False) +class Object: + """some random weakreffable object""" + pass + +def test_weakref_ref_with_callback(): + obj = Object() + r = ref(obj, lambda x: x) + assert r() is obj + obj = None + assert r() is None, r() diff --git a/mypyc/test/test_irbuild.py b/mypyc/test/test_irbuild.py index 9c0ad06416a7..d8f974ef201b 100644 --- a/mypyc/test/test_irbuild.py +++ b/mypyc/test/test_irbuild.py @@ -53,6 +53,7 @@ "irbuild-constant-fold.test", "irbuild-glue-methods.test", "irbuild-math.test", + "irbuild-weakref.test", ] if sys.version_info >= (3, 10): diff --git a/mypyc/test/test_run.py b/mypyc/test/test_run.py index e5b7e2421433..cd32daee4bdc 100644 --- a/mypyc/test/test_run.py +++ b/mypyc/test/test_run.py @@ -69,6 +69,7 @@ "run-dunders-special.test", "run-singledispatch.test", "run-attrs.test", + "run-weakref.test", "run-python37.test", "run-python38.test", ]