diff --git a/mypyc/lib-rt/misc_ops.c b/mypyc/lib-rt/misc_ops.c index 870cd4682ea7..9aca9c85e17b 100644 --- a/mypyc/lib-rt/misc_ops.c +++ b/mypyc/lib-rt/misc_ops.c @@ -950,8 +950,9 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, PyObject *typing = NULL; PyObject *get_type_hints = NULL; PyObject *type_hints = NULL; + PyObject *result = NULL; - if (registry == NULL) goto fail; + if (registry == NULL) goto done; if (func == NULL) { // one argument case if (PyType_Check(cls)) { @@ -959,8 +960,9 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, // bind cls to the first argument so that register gets called again with both the // class and the function register_func = PyObject_GetAttr(singledispatch_func, mypyc_interned_str.register_); - if (register_func == NULL) goto fail; - return PyMethod_New(register_func, cls); + if (register_func == NULL) goto done; + result = PyMethod_New(register_func, cls); + goto done; } // passed a function PyObject *annotations = PyFunction_GetAnnotations(cls); @@ -971,49 +973,49 @@ PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, if (annotations == NULL) { PyErr_Format(PyExc_TypeError, invalid_first_arg_msg, cls); - goto fail; + goto done; } - Py_INCREF(annotations); - func = cls; typing = PyImport_ImportModule("typing"); - if (typing == NULL) goto fail; + if (typing == NULL) goto done; get_type_hints = PyObject_GetAttr(typing, mypyc_interned_str.get_type_hints); + if (get_type_hints == NULL) goto done; type_hints = PyObject_CallOneArg(get_type_hints, func); + if (type_hints == NULL) goto done; PyObject *argname; Py_ssize_t pos = 0; if (!PyDict_Next(type_hints, &pos, &argname, &cls)) { // the functools implementation raises the same type error if annotations is an empty dict PyErr_Format(PyExc_TypeError, invalid_first_arg_msg, cls); - goto fail; + goto done; } if (!PyType_Check(cls)) { const char *invalid_annotation_msg = "Invalid annotation for %R. %R is not a class."; PyErr_Format(PyExc_TypeError, invalid_annotation_msg, argname, cls); - goto fail; + goto done; } } if (PyDict_SetItem(registry, cls, func) == -1) { - goto fail; + goto done; } // clear the cache so we consider the newly added function when dispatching PyObject *dispatch_cache = PyObject_GetAttr(singledispatch_func, mypyc_interned_str.dispatch_cache); - if (dispatch_cache == NULL) goto fail; + if (dispatch_cache == NULL) goto done; PyDict_Clear(dispatch_cache); + Py_DECREF(dispatch_cache); - Py_INCREF(func); - return func; + result = Py_NewRef(func); -fail: +done: Py_XDECREF(registry); Py_XDECREF(register_func); Py_XDECREF(typing); Py_XDECREF(get_type_hints); Py_XDECREF(type_hints); - return NULL; + return result; } diff --git a/mypyc/test-data/run-singledispatch.test b/mypyc/test-data/run-singledispatch.test index 03b937261e22..973a601ba855 100644 --- a/mypyc/test-data/run-singledispatch.test +++ b/mypyc/test-data/run-singledispatch.test @@ -724,3 +724,85 @@ assert f(C()) == 'c' [file driver.py] import register + +[case testRegisterAnnotationErrorOwnership] +import sys +from functools import singledispatch +from typing import Any +from testutil import assertRaises +from helpers import invalid + +@singledispatch +def fun(arg: Any) -> bool: + return False + +def test_annotation_error_ownership() -> None: + getrefcount = getattr(sys, "getrefcount") + annotations = getattr(invalid, "__annotations__") + before = getrefcount(annotations) + for _ in range(5): + with assertRaises(NameError): + fun.register(invalid) + assert getrefcount(annotations) == before + +[file helpers.py] +# Annotation-based registration requires an interpreted Python function. +def invalid(arg: "UnknownType") -> bool: # type: ignore[name-defined] + return False + +[case testRegisterSuccessOwnership] +import gc +import sys +import weakref +from functools import singledispatch +from mypy_extensions import mypyc_attr +from typing import Any +from helpers import specialized + +@singledispatch +def fun(arg: Any) -> bool: + return False + +class Arg: + pass + +# Keep the dict subclass weak-referenceable. +@mypyc_attr(native_class=False) +class Hints(dict): + pass + +def test_success_ownership() -> None: + refs: list[weakref.ReferenceType[Hints]] = [] + def get_type_hints(func: Any) -> Hints: + hints = Hints(arg=Arg) + refs.append(weakref.ref(hints)) + return hints + + getrefcount = getattr(sys, "getrefcount") + registry = getattr(fun, "registry") + cache = getattr(fun, "dispatch_cache") + registry_before = getrefcount(registry) + cache_before = getrefcount(cache) + # Access typing dynamically to work with the minimal test stubs. + typing = getattr(sys, "modules")["typing"] + original_get_type_hints = typing.get_type_hints + typing.get_type_hints = get_type_hints + try: + for _ in range(5): + fun.register(specialized) + assert fun(Arg()) + fun.register(int)(specialized) + assert fun(1) + fun.register(str, specialized) + assert fun("") + finally: + typing.get_type_hints = original_get_type_hints + gc.collect() + assert all(ref() is None for ref in refs) + assert getrefcount(registry) == registry_before + assert getrefcount(cache) == cache_before + +[file helpers.py] +# Annotation-based registration requires an interpreted Python function. +def specialized(arg: object) -> bool: + return True