[mypyc] Fix crash on multiple nested decorated functions with same name (#20666)

This happens when a nested function has different definitions in
different condition branches. mypyc currently crashes in this case when
calling `get_func_target` in `transform_decorator`
https://github.com/python/mypy/blob/master/mypyc/irbuild/function.py#L118.

The `FuncDef` objects for the definitions after the first one will have
the `original_def` attribute set. This makes `get_func_target` lookup
this original def but it doesn't find it when the original def is a
decorator because registers are only created for `FuncDef`s.

Extracting a `FuncDef` from the original decorator and looking it up
instead fixes the crash.

Author: p-sawicki

mypyc/irbuild/function.py

@@ -852,9 +852,11 @@ def get_func_target(builder: IRBuilder, fdef: FuncDef) -> AssignmentTarget:
     If the function was not already defined somewhere, then define it
     and add it to the current environment.
     """
-    if fdef.original_def:
+    if orig := fdef.original_def:
+        if isinstance(orig, Decorator):
+            orig = orig.func
         # Get the target associated with the previously defined FuncDef.
-        return builder.lookup(fdef.original_def)
+        return builder.lookup(orig)
 
     if builder.fn_info.is_generator or builder.fn_info.add_nested_funcs_to_env:
         return builder.lookup(fdef)

mypyc/test-data/run-async.test

@@ -1809,7 +1809,7 @@ from typing import Any, Callable, TypeVar, cast
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-def mult(x: int) -> Callable[[F], F]:
+def mult_different_wrapper_names(x: int) -> Callable[[F], F]:
     def decorate(fn: F) -> F:
         def get_multiplier() -> int:
             return x
@@ -1829,18 +1829,47 @@ def mult(x: int) -> Callable[[F], F]:
 
     return decorate
 
-@mult(3)
+def mult_same_wrapper_names(x: int) -> Callable[[F], F]:
+    def decorate(fn: F) -> F:
+        def get_multiplier() -> int:
+            return x
+
+        if inspect.iscoroutinefunction(fn):
+            @functools.wraps(fn)
+            async def wrapper(*args, **kwargs) -> Any:
+                return get_multiplier() * await fn(*args, **kwargs)
+        else:
+            @functools.wraps(fn)
+            def wrapper(*args, **kwargs) -> Any:
+                return get_multiplier() * fn(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    return decorate
+
+@mult_different_wrapper_names(3)
 def identity(x: int):
     return x
 
-@mult(5)
+@mult_different_wrapper_names(5)
 async def async_identity(x: int):
     return x
 
+@mult_same_wrapper_names(2)
+def times_two(x: int):
+    return x * 2
+
+@mult_same_wrapper_names(4)
+async def async_times_two(x: int):
+    return x * 2
+
 def test_nested_coroutine_calls_another_nested_function():
     assert identity(1) == 3
     assert asyncio.run(async_identity(2)) == 10
 
+    assert times_two(3) == 12
+    assert asyncio.run(async_times_two(4)) == 32
+
 [file asyncio/__init__.pyi]
 from typing import Any, Generator
 

mypyc/test-data/run-functions.test

@@ -154,6 +154,62 @@ def if_else(flag: int) -> str:
             return 'if_else.inner: third definition'
     return inner()
 
+def wrap(f: Callable[[], str]):
+    def inner():
+        return 'wrapped ' + f()
+    return inner
+
+def if_else_all_decorated(flag: int) -> str:
+    def dummy_function() -> str:
+        return 'if_else.dummy_function'
+
+    if flag < 0:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: first definition'
+    elif flag > 0:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: second definition'
+    else:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: third definition'
+    return inner()
+
+def if_else_first_decorated(flag: int) -> str:
+    def dummy_function() -> str:
+        return 'if_else.dummy_function'
+
+    if flag < 0:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: first definition'
+    elif flag > 0:
+        def inner() -> str:
+            return 'if_else.inner: second definition'
+    else:
+        def inner() -> str:
+            return 'if_else.inner: third definition'
+    return inner()
+
+def if_else_all_but_first_decorated(flag: int) -> str:
+    def dummy_function() -> str:
+        return 'if_else.dummy_function'
+
+    if flag < 0:
+        def inner() -> str:
+            return 'if_else.inner: first definition'
+    elif flag > 0:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: second definition'
+    else:
+        @wrap
+        def inner() -> str:
+            return 'if_else.inner: third definition'
+    return inner()
+
 def for_loop() -> int:
     def dummy_function() -> str:
         return 'for_loop.dummy_function'
@@ -235,8 +291,9 @@ toplevel_lambda = lambda x: 10 + global_upvar + x
 
 [file driver.py]
 from native import (
-    a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, triple, if_else, for_loop, while_loop,
-    free_vars, lambdas, outer, inner, A, toplevel_lambda
+    a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, triple, if_else, if_else_all_decorated,
+    if_else_first_decorated, if_else_all_but_first_decorated,
+    for_loop, while_loop, free_vars, lambdas, outer, inner, A, toplevel_lambda
 )
 
 assert a()() == None
@@ -266,6 +323,18 @@ assert if_else(-1) == 'if_else.inner: first definition'
 assert if_else(1) == 'if_else.inner: second definition'
 assert if_else(0) == 'if_else.inner: third definition'
 
+assert if_else_all_decorated(-1) == 'wrapped if_else.inner: first definition'
+assert if_else_all_decorated(1) == 'wrapped if_else.inner: second definition'
+assert if_else_all_decorated(0) == 'wrapped if_else.inner: third definition'
+
+assert if_else_first_decorated(-1) == 'wrapped if_else.inner: first definition'
+assert if_else_first_decorated(1) == 'if_else.inner: second definition'
+assert if_else_first_decorated(0) == 'if_else.inner: third definition'
+
+assert if_else_all_but_first_decorated(-1) == 'if_else.inner: first definition'
+assert if_else_all_but_first_decorated(1) == 'wrapped if_else.inner: second definition'
+assert if_else_all_but_first_decorated(0) == 'wrapped if_else.inner: third definition'
+
 assert for_loop() == 3
 assert while_loop() == 3