[mypyc] Fix reference leaks of temporaries in singledispatch registration

mypyc/lib-rt/misc_ops.c

@@ -950,17 +950,19 @@ 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)) {
             // passed a class
             // 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;
 
 }
 

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