[mypyc] Fix crash on double yielding Iterators (#21826)

Closes mypyc/mypyc#1210

This fixes another issue I ran into while trying to use `mypyc` for
`isort`.

I am aware the contributing guidelines say new contributors are not
encouraged to use LLMs but I hope this can get the same handling as
#21785. I have tried to ensure this
patch works as much as I can by:

Testing locally with `pytest mypyc/test/test_run.py `.
Tested that the test fails on `master` without the changes.
Also tried to run CI on my local fork
(DanielNoord#2), which succeeded.

As with the previous PR, feel free to push changes to the branch or
cherry pick this into another branch. I just want to unblock `isort` :)

Author: DanielNoord

mypyc/irbuild/generator.py

@@ -166,6 +166,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

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)