From 04903e87d3df082ae5dce0287b68ade0602217c5 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Wed, 7 Oct 2026 15:18:11 +0100
Subject: [PATCH] [mypyc] Fix crash on nested function in a method of a
 non-native class

Fixes mypyc/mypyc#1225. Fixes mypyc/mypyc#1205.

Compiling a method of a non-native class that contains a nested function
or a lambda crashed with an internal error:

```python
from mypy_extensions import mypyc_attr

@mypyc_attr(native_class=False)
class C:
    def f(self) -> Callable[[], C]:
        return lambda: self  # AssertionError
```

Methods of non-native classes are compiled to callable classes, like
nested functions. If a method contains a nested function or a lambda,
the method creates an environment object, which holds the variables
captured by the nested function (`self` above) and the nested function
objects. When adding these to the environment, the IR builder treated
the method like a nested function and looked for the environment via
the callable class, but only nested functions set it up there. Now the
method's own environment is used, as in other functions that aren't
nested.
---
 mypyc/irbuild/env_class.py       |  2 +-
 mypyc/test-data/run-classes.test | 94 ++++++++++++++++++++++++--------
 2 files changed, 73 insertions(+), 23 deletions(-)

diff --git a/mypyc/irbuild/env_class.py b/mypyc/irbuild/env_class.py
index 2d91d1f32591..bce61c3fa791 100644
--- a/mypyc/irbuild/env_class.py
+++ b/mypyc/irbuild/env_class.py
@@ -293,7 +293,7 @@ def add_vars_to_env(builder: IRBuilder, prefix: str = "") -> None:
     env_for_func: FuncInfo | ImplicitClass = fn_info
     if fn_info.is_generator:
         env_for_func = fn_info.generator_class
-    elif (fn_info.is_nested or fn_info.in_non_ext) and not fn_info.is_comprehension_scope:
+    elif fn_info.is_nested and not fn_info.is_comprehension_scope:
         env_for_func = fn_info.callable_class
 
     if fn_info.fitem in builder.free_variables:
diff --git a/mypyc/test-data/run-classes.test b/mypyc/test-data/run-classes.test
index 9b204f1da3b6..9c50e249ea3b 100644
--- a/mypyc/test-data/run-classes.test
+++ b/mypyc/test-data/run-classes.test
@@ -259,7 +259,9 @@ def test_non_ext_any() -> None:
         n.y
 
 [case testNonExtMisc]
-from typing import Any, overload
+from typing import Any, Callable, Iterator, overload
+
+from mypy_extensions import mypyc_attr
 
 def decorator(cls) -> Any:
     return cls
@@ -305,32 +307,80 @@ class Overload:
     def get(self, index: Any) -> Any:
         return index
 
-def get(c: Overload, s: str) -> str:
-    return c.get(s)
-
 @decorator
 class Var:
     x = 'xy'
 
-def get_class_var() -> str:
-    return Var.x
+def call(f: Callable[[int], int], x: int) -> int:
+    return f(x)
 
-[file driver.py]
-from native import A, Overload, get, get_class_var
-a = A()
-assert a.a == 1
-assert a.b == 2
-assert a.c == 3
-assert a.get_a() == 1
-assert a.get_b() == 2
-assert a.get_c() == 3
-assert A.constant() == 4
-
-o = Overload()
-assert get(o, "test") == "test"
-assert o.get(20) == 20
-
-assert get_class_var() == 'xy'
+@mypyc_attr(native_class=False)
+class NestedFunc:
+    def __init__(self, n: int) -> None:
+        self.n = n
+
+    def no_capture(self, x: int) -> int:
+        def f(y: int) -> int:
+            return y + 1
+        return call(f, x)
+
+    def conditional(self, x: int) -> int:
+        if x:
+            def f(y: int) -> int:
+                return y * 2
+            return call(f, x)
+        return -1
+
+    def capture(self, x: int) -> int:
+        def f(y: int) -> int:
+            return y + x + self.n
+        return call(f, 1)
+
+    def recursive(self, x: int) -> int:
+        def f(y: int) -> int:
+            return 0 if y == 0 else y + f(y - 1)
+        return f(x)
+
+    def with_lambda(self, x: int) -> int:
+        return call(lambda y: y + x + self.n, 1)
+
+    def return_lambda(self) -> Callable[[], "NestedFunc"]:
+        return lambda: self
+
+    def gen(self, x: int) -> Iterator[int]:
+        def f(y: int) -> int:
+            return y + x
+        yield call(f, 1)
+        yield call(f, 2)
+
+def test_inheritance() -> None:
+    a = A()
+    assert a.a == 1
+    assert a.b == 2
+    assert a.c == 3
+    assert a.get_a() == 1
+    assert a.get_b() == 2
+    assert a.get_c() == 3
+    assert A.constant() == 4
+
+def test_overload() -> None:
+    o = Overload()
+    assert o.get("test") == "test"
+    assert o.get(20) == 20
+
+def test_class_var() -> None:
+    assert Var.x == 'xy'
+
+def test_nested_functions_in_methods() -> None:
+    n = NestedFunc(10)
+    assert n.no_capture(1) == 2
+    assert n.conditional(3) == 6
+    assert n.conditional(0) == -1
+    assert n.capture(2) == 13
+    assert n.recursive(4) == 10
+    assert n.with_lambda(2) == 13
+    assert n.return_lambda()() is n
+    assert list(n.gen(5)) == [6, 7]
 
 [case testNonExtClassDecoratorReturningNonClass]
 from functools import cache