From 009644dc8778b87c8b766a8314283940bdd3de59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20van=20Noord?=
 <13665637+DanielNoord@users.noreply.github.com>
Date: Sun, 9 Aug 2026 22:53:35 +0200
Subject: [PATCH] Fix crash on double yielding Iterators

---
 mypyc/irbuild/generator.py          |  6 ++++++
 mypyc/test-data/run-generators.test | 17 ++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/mypyc/irbuild/generator.py b/mypyc/irbuild/generator.py
index 3e7f18c91c753..555baaada0e81 100644
--- a/mypyc/irbuild/generator.py
+++ b/mypyc/irbuild/generator.py
@@ -174,6 +174,12 @@ def setup_generator_class(builder: IRBuilder) -> ClassIR:
         builder.fn_info.env_class = generator_class_ir
     else:
         generator_class_ir.attributes[ENV_ATTR_NAME] = RInstance(builder.fn_info.env_class)
+        if not builder.fn_info.fitem.is_coroutine:
+            # After completion generators still need generator.__mypyc_env__ for subsequent
+            # __next__() calls to observe the terminal next-label and raise StopIteration.
+            # Coroutines can't be resumed after completion, so keeping the environment alive
+            # there would just extend local lifetimes unnecessarily.
+            generator_class_ir.attrs_to_keep_alive_on_completion.add(ENV_ATTR_NAME)
 
     builder.classes.append(generator_class_ir)
     return generator_class_ir
diff --git a/mypyc/test-data/run-generators.test b/mypyc/test-data/run-generators.test
index a23af2ed6eea7..f4cef93ea8f3e 100644
--- a/mypyc/test-data/run-generators.test
+++ b/mypyc/test-data/run-generators.test
@@ -1,7 +1,7 @@
 # Test cases for generators and yield (compile and run)
 
 [case testYield]
-from typing import Generator, Iterable, Union, Tuple, Dict
+from typing import Callable, Dict, Generator, Iterable, Iterator, Tuple, Union
 
 def yield_three_times() -> Iterable[int]:
     yield 1
@@ -87,6 +87,19 @@ def return_tuple() -> Generator[int, None, Tuple[int, int]]:
     yield 0
     return 1, 2
 
+def call_lambda(fn: Callable[[], None]) -> None:
+    fn()
+
+def get_iterator() -> Iterator[str]:
+    call_lambda(lambda: None)
+    yield ""
+
+def yield_twice_via_generator_reuse() -> Iterator[str]:
+    identified_imports = get_iterator()
+    yield from identified_imports
+    for identified_import in identified_imports:
+        yield identified_import
+
 [file driver.py]
 from native import (
     yield_three_times,
@@ -99,6 +112,7 @@ from native import (
     A,
     return_tuple,
     yield_dict_methods,
+    yield_twice_via_generator_reuse,
 )
 from testutil import run_generator
 from collections import defaultdict
@@ -114,6 +128,7 @@ assert run_generator(A(0).generator()) == ((0,), None)
 assert run_generator(return_tuple()) == ((0,), (1, 2))
 assert run_generator(yield_dict_methods({}, {}, {})) == ((), None)
 assert run_generator(yield_dict_methods({1: 2}, {3: 4}, {5: 6})) == ((1, 3, 4, 6), None)
+assert run_generator(yield_twice_via_generator_reuse()) == (("",), None)
 dd = defaultdict(int, {0: 1})
 assert run_generator(yield_dict_methods(dd, dd, dd)) == ((0, 0, 1, 1), None)