# Test cases for generators and yield (compile and run)

[case testYield]
from typing import Generator, Iterable, Union, Tuple, Dict

def yield_three_times() -> Iterable[int]:
    yield 1
    yield 2
    yield 3

def yield_twice_and_return() -> Generator[int, None, int]:
    yield 1
    yield 2
    return 4

def yield_while_loop() -> Generator[int, None, int]:
    i = 0
    while i < 5:
        if i == 3:
            return i
        yield i
        i += 1
    return -1

def yield_for_loop() -> Iterable[int]:
    l = [i for i in range(3)]
    for i in l:
        yield i

    d = {k: None for k in range(3)}
    for k in d:
        yield k

    for i in range(3):
        yield i

    for i in range(three()):
        yield i

def yield_with_except() -> Generator[int, None, None]:
    yield 10
    try:
        return
    except:
        print('Caught exception inside generator function')

def complex_yield(a: int, b: str, c: float) -> Generator[Union[str, int], None, float]:
    x = 2
    while x < a:
        if x % 2 == 0:
            dummy_var = 1
            yield str(x) + ' ' + b
            dummy_var = 1
        else:
            dummy_var = 1
            yield x
            dummy_var = 1
        x += 1
    return c

def yield_with_default(x: bool = False) -> Iterable[int]:
    if x:
        yield 0

def yield_dict_methods(d1: Dict[int, int],
                       d2: Dict[int, int],
                       d3: Dict[int, int]) -> Iterable[int]:
    for k in d1.keys():
        yield k
    for k, v in d2.items():
        yield k
        yield v
    for v in d3.values():
        yield v

def three() -> int:
    return 3

class A(object):
    def __init__(self, x: int) -> None:
        self.x = x

    def generator(self) -> Iterable[int]:
        yield self.x

def return_tuple() -> Generator[int, None, Tuple[int, int]]:
    yield 0
    return 1, 2

[file driver.py]
from native import (
    yield_three_times,
    yield_twice_and_return,
    yield_while_loop,
    yield_for_loop,
    yield_with_except,
    complex_yield,
    yield_with_default,
    A,
    return_tuple,
    yield_dict_methods,
)
from testutil import run_generator
from collections import defaultdict

assert run_generator(yield_three_times()) == ((1, 2, 3), None)
assert run_generator(yield_twice_and_return()) == ((1, 2), 4)
assert run_generator(yield_while_loop()) == ((0, 1, 2), 3)
assert run_generator(yield_for_loop()) == (tuple(4 * [i for i in range(3)]), None)
assert run_generator(yield_with_except()) == ((10,), None)
assert run_generator(complex_yield(5, 'foo', 1.0)) == (('2 foo', 3, '4 foo'), 1.0)
assert run_generator(yield_with_default()) == ((), None)
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)
dd = defaultdict(int, {0: 1})
assert run_generator(yield_dict_methods(dd, dd, dd)) == ((0, 0, 1, 1), None)

for i in yield_twice_and_return():
    print(i)

for i in yield_while_loop():
    print(i)

[out]
1
2
0
1
2

[case testYieldTryFinallyWith]
from typing import Generator, Any

class Thing:
    def __init__(self, x: str) -> None:
        self.x = x
    def __enter__(self) -> str:
        print('enter!', self.x)
        if self.x == 'crash':
            raise Exception('ohno')
        return self.x
    def __exit__(self, x: Any, y: Any, z: Any) -> None:
        print('exit!', self.x, y)

def yield_try_finally() -> Generator[int, None, str]:
    try:
        yield 1
        yield 2
        return 'lol'
    except Exception:
        raise
    finally:
        print('goodbye!')

def yield_with(i: int) -> Generator[int, None, int]:
    with Thing('a') as x:
        yield 1
        print("yooo?", x)
        if i == 0:
            yield 2
            return 10
        elif i == 1:
            raise Exception('exception!')
    return -1

[file driver.py]
from native import yield_try_finally, yield_with
from testutil import run_generator

print(run_generator(yield_try_finally(), p=True))
print(run_generator(yield_with(0), p=True))
print(run_generator(yield_with(1), p=True))
[out]
1
2
goodbye!
((1, 2), 'lol')
enter! a
1
yooo? a
2
exit! a None
((1, 2), 10)
enter! a
1
yooo? a
exit! a exception!
((1,), 'exception!')

[case testYieldNested]
from typing import Callable, Generator, Iterator, TypeVar, overload

from testutil import run_generator

def normal(a: int, b: float) -> Callable:
    def generator(x: int, y: str) -> Generator:
        yield a
        yield b
        yield x
        yield y
    return generator

def generator(a: int) -> Generator:
    def normal(x: int) -> int:
        return a + x
    for i in range(3):
        yield normal(i)

def triple() -> Callable:
    def generator() -> Generator:
        x = 0
        def inner() -> int:
            x += 1
            return x
        while x < 3:
            yield inner()
    return generator

def another_triple() -> Callable:
    def generator() -> Generator:
        x = 0
        def inner_generator() -> Generator:
            x += 1
            yield x
        yield next(inner_generator())
    return generator

def outer() -> Generator:
    def recursive(n: int) -> Generator:
        if n < 10:
            for i in range(n):
                yield i
            return
        for i in recursive(5):
            yield i
    return recursive(10)

def test_return_nested_generator() -> None:
    assert run_generator(normal(1, 2.0)(3, '4.00')) == ((1, 2.0, 3, '4.00'), None)
    assert run_generator(generator(1)) == ((1, 2, 3), None)
    assert run_generator(triple()()) == ((1, 2, 3), None)
    assert run_generator(another_triple()()) == ((1,), None)
    assert run_generator(outer()) == ((0, 1, 2, 3, 4), None)

def call_nested(x: int) -> list[int]:
    def generator() -> Iterator[int]:
        n = int() + 2
        yield x
        yield n * x

    a = []
    for x in generator():
        a.append(x)
    return a

T = TypeVar("T")

def deco(f: T) -> T:
    return f

def call_nested_decorated(x: int) -> list[int]:
    @deco
    def generator() -> Iterator[int]:
        n = int() + 3
        yield x
        yield n * x

    a = []
    for x in generator():
        a.append(x)
    return a

def call_nested_recursive(x: int) -> Iterator:
    def recursive(x: int) -> Iterator:
        if x > 0:
            yield from recursive(x - 1)
        yield x

    yield from recursive(x)

def test_call_nested_generator_in_function() -> None:
    assert call_nested_decorated(5) == [5, 15]
    assert list(call_nested_recursive(5)) == [0, 1, 2, 3, 4, 5]

[case testYieldThrow]
from typing import Generator, Iterable, Any, Union
from traceback import print_tb
from contextlib import contextmanager
import wrapsys

def generator() -> Generator[int, None, Union[int, None]]:
    try:
        yield 1
        yield 2
        yield 3
    except Exception as e:
        print_tb(wrapsys.exc_info()[2])
        s = str(e)
        if s:
            print('caught exception with value ' + s)
        else:
            print('caught exception without value')
        return 0
    return None

def no_except() -> Iterable[int]:
    yield 1
    yield 2

def raise_something() -> Iterable[int]:
    yield 1
    yield 2
    raise Exception('failure')

def wrapper(x: Any) -> Any:
    return (yield from x)

def foo() -> Generator[int, None, None]:
    try:
        yield 1
    except Exception as e:
        print(str(e))
    finally:
        print('goodbye')

ctx_manager = contextmanager(foo)

[file wrapsys.py]
# This is a gross hack around some limitations of the test system/mypyc.
from typing import Any
import sys
def exc_info() -> Any:
    return sys.exc_info()  # type: ignore

[file driver.py]
import sys
from typing import Generator, Tuple, TypeVar, Sequence
from native import generator, ctx_manager, wrapper, no_except, raise_something

T = TypeVar('T')
U = TypeVar('U')

def run_generator_and_throw(gen: Generator[T, None, U],
                            num_times: int,
                            value: object = None,
                            traceback: object = None) -> Tuple[Sequence[T], U]:
    res = []
    try:
        for i in range(num_times):
            res.append(next(gen))
        if value is not None and traceback is not None:
            gen.throw(Exception, value, traceback)
        elif value is not None:
            gen.throw(Exception, value)
        else:
            gen.throw(Exception)
    except StopIteration as e:
        return (tuple(res), e.value)
    except Exception as e:
        return (tuple(res), str(e))

assert run_generator_and_throw(generator(), 0, 'hello') == ((), 'hello')
assert run_generator_and_throw(generator(), 3) == ((1, 2, 3), 0)
assert run_generator_and_throw(generator(), 2, 'some string') == ((1, 2), 0)
try:
    raise Exception
except Exception as e:
    tb = sys.exc_info()[2]
    assert run_generator_and_throw(generator(), 1, 'some other string', tb) == ((1,), 0)

assert run_generator_and_throw(wrapper(generator()), 0, 'hello') == ((), 'hello')
assert run_generator_and_throw(wrapper(generator()), 3) == ((1, 2, 3), 0)
assert run_generator_and_throw(wrapper(generator()), 2, 'some string') == ((1, 2), 0)
# Make sure we aren't leaking exc_info
assert sys.exc_info()[0] is None

assert run_generator_and_throw(wrapper([1, 2, 3]), 3, 'lol') == ((1, 2, 3), 'lol')
assert run_generator_and_throw(wrapper(no_except()), 2, 'lol') == ((1, 2), 'lol')

assert run_generator_and_throw(wrapper(raise_something()), 3) == ((1, 2), 'failure')

with ctx_manager() as c:
    raise Exception('exception')

[out]
  File "native.py", line 10, in generator
    yield 3
  File "native.py", line 9, in generator
    yield 2
  File "native.py", line 8, in generator
    yield 1
  File "driver.py", line 31, in <module>
    raise Exception
  File "native.py", line 10, in generator
    yield 3
  File "native.py", line 31, in wrapper
    return (yield from x)
  File "native.py", line 9, in generator
    yield 2
  File "native.py", line 31, in wrapper
    return (yield from x)
caught exception without value
caught exception with value some string
caught exception with value some other string
caught exception without value
caught exception with value some string
exception
goodbye

[case testYieldSend]
from typing import Generator

def basic() -> Generator[int, int, int]:
    x = yield 1
    y = yield (x + 1)
    return y

def use_from() -> Generator[int, int, int]:
    return (yield from basic())

[file driver.py]
from native import basic, use_from
from testutil import run_generator

assert run_generator(basic(), [5, 50]) == ((1, 6), 50)
assert run_generator(use_from(), [5, 50]) == ((1, 6), 50)

[case testYieldFrom]
from typing import Generator, Iterator, List

def basic() -> Iterator[int]:
    yield from [1, 2, 3]

def call_next() -> int:
    x = []  # type: List[int]
    return next(iter(x))

def inner(b: bool) -> Generator[int, None, int]:
    if b:
        yield from [1, 2, 3]
    return 10

def with_return(b: bool) -> Generator[int, None, int]:
    x = yield from inner(b)
    for a in [1, 2]:
        pass
    return x

[file driver.py]
from native import basic, call_next, with_return
from testutil import run_generator, assertRaises

assert run_generator(basic()) == ((1, 2, 3), None)

with assertRaises(StopIteration):
    call_next()

assert run_generator(with_return(True)) == ((1, 2, 3), 10)
assert run_generator(with_return(False)) == ((), 10)

[case testNextGenerator]
from typing import Iterable

def f(x: int) -> int:
    print(x)
    return x

def call_next_loud(l: Iterable[int], val: int) -> int:
    return next(i for i in l if f(i) == val)

def call_next_default(l: Iterable[int], val: int) -> int:
    return next((i*2 for i in l if i == val), -1)

def call_next_default_list(l: Iterable[int], val: int) -> int:
    return next((i*2 for i in l if i == val), -1)
[file driver.py]
from native import call_next_loud, call_next_default, call_next_default_list
from testutil import assertRaises

assert call_next_default([0, 1, 2], 0) == 0
assert call_next_default([0, 1, 2], 1) == 2
assert call_next_default([0, 1, 2], 2) == 4
assert call_next_default([0, 1, 2], 3) == -1
assert call_next_default([], 0) == -1
assert call_next_default_list([0, 1, 2], 0) == 0
assert call_next_default_list([0, 1, 2], 1) == 2
assert call_next_default_list([0, 1, 2], 2) == 4
assert call_next_default_list([0, 1, 2], 3) == -1
assert call_next_default_list([], 0) == -1

assert call_next_loud([0, 1, 2], 0) == 0
assert call_next_loud([0, 1, 2], 1) == 1
assert call_next_loud([0, 1, 2], 2) == 2
with assertRaises(StopIteration):
    call_next_loud([42], 3)
with assertRaises(StopIteration):
    call_next_loud([], 3)

[out]
0
0
1
0
1
2
42

[case testGeneratorSuper]
from typing import Iterator, Callable, Any

class A():
    def testA(self) -> int:
        return 2

class B(A):
    def testB(self) -> Iterator[int]:
        x = super().testA()
        while True:
            yield x

def testAsserts():
    b = B()
    b_gen = b.testB()
    assert next(b_gen) == 2

[file driver.py]
from native import testAsserts

testAsserts()

[case testNameClashIssues]
class A:
    def foo(self) -> object:
        yield
class B:
    def foo(self) -> object:
        yield

class C:
    def foo(self) -> None:
        def bar(self) -> None:
            pass

def C___foo() -> None: pass

class D:
    def foo(self) -> None:
        def bar(self) -> None:
            pass

class E:
    default: int
    switch: int

[file driver.py]
# really I only care it builds

[case testCloseStopIterationRaised]
def g() -> object:
    try:
        yield 1
    except GeneratorExit:
        raise

[file driver.py]
from native import g

gen = g()
next(gen)
gen.close()

[case testCloseGeneratorExitRaised]
def g() -> object:
    yield 1

[file driver.py]
from native import g

gen = g()
next(gen)
gen.close()

[case testCloseGeneratorExitIgnored]
def g() -> object:
    try:
        yield 1
    except GeneratorExit:
        pass

    yield 2

[file driver.py]
from native import g

gen = g()
next(gen)
try:
    gen.close()
except RuntimeError as e:
    assert str(e) == 'generator ignored GeneratorExit'
else:
    assert False

[case testCloseGeneratorRaisesAnotherException]
def g() -> object:
    try:
        yield 1
    except GeneratorExit:
        raise RuntimeError("error")

[file driver.py]
from native import g

gen = g()
next(gen)
try:
    gen.close()
except RuntimeError as e:
    assert str(e) == 'error'
else:
    assert False

[case testBorrowingInGeneratorNearYield]
from typing import Iterator

class Foo:
    flag = False

class C:
    foo = Foo()

    def genf(self) -> Iterator[None]:
        self.foo.flag = True
        yield
        self.foo.flag = False

def test_near_yield() -> None:
    c = C()
    for x in c.genf():
        pass
    assert c.foo.flag == False

[case testGeneratorEarlyReturnWithBorrows]
from typing import Iterator
class Bar:
    bar = 0
class Foo:
    bar = Bar()
    def f(self) -> Iterator[int]:
        if self:
            self.bar.bar += 1
            return
        yield 0

def test_early_return() -> None:
    foo = Foo()
    for x in foo.f():
        pass
    assert foo.bar.bar == 1

[case testBorrowingInGeneratorInTupleAssignment]
from typing import Iterator

class Foo:
    flag1: bool
    flag2: bool

class C:
    foo: Foo

    def genf(self) -> Iterator[None]:
        self.foo.flag1, self.foo.flag2 = True, True
        yield
        self.foo.flag1, self.foo.flag2 = False, False

def test_generator() -> None:
    c = C()
    c.foo = Foo()
    gen = c.genf()
    next(gen)
    assert c.foo.flag1 == c.foo.flag2 == True
    assert list(gen) == []
    assert c.foo.flag1 == c.foo.flag2 == False


[case testYieldInFinally]
from typing import Generator

def finally_yield() -> Generator[str, None, str]:
    try:
        return 'test'
    finally:
        yield 'x'


[file driver.py]
from native import finally_yield
from testutil import run_generator

yields, val = run_generator(finally_yield())
assert yields == ('x',)
assert val == 'test', val

[case testUnreachableComprehensionNoCrash]
from typing import List

def list_comp() -> List[int]:
    if True:
        return [5]
    return [i for i in [5]]

[file driver.py]
from native import list_comp
assert list_comp() == [5]

[case testWithNative]
class DummyContext:
    def __init__(self) -> None:
        self.x = 0

    def __enter__(self) -> None:
        self.x += 1

    def __exit__(self, exc_type, exc_value, exc_tb) -> None:
        self.x -= 1

def test_basic() -> None:
    context = DummyContext()
    with context:
        assert context.x == 1
    assert context.x == 0

[case testYieldSpill]
from typing import Generator
from testutil import run_generator

def f() -> int:
    return 1

def yield_spill() -> Generator[str, int, int]:
    return f() + (yield "foo")

def test_basic() -> None:
    x = run_generator(yield_spill(), [2])
    yields, val = x
    assert yields == ('foo',)
    assert val == 3, val

[case testGeneratorReuse]
from typing import Iterator, Any

def gen(x: list[int]) -> Iterator[list[int]]:
    y = [9]
    for z in x:
        yield y + [z]
    yield y

def gen_range(n: int) -> Iterator[int]:
    for x in range(n):
        yield x

def test_use_generator_multiple_times_one_at_a_time() -> None:
    for i in range(100):
        a = []
        for x in gen([2, i]):
            a.append(x)
        assert a == [[9, 2], [9, i], [9]]

def test_use_multiple_generator_instances_at_same_time() -> None:
    a = []
    for x in gen([2]):
        a.append(x)
        for y in gen([3, 4]):
            a.append(y)
    assert a == [[9, 2], [9, 3], [9, 4], [9], [9], [9, 3], [9, 4], [9]]

def test_use_multiple_generator_instances_at_same_time_2() -> None:
    a = []
    for x in gen_range(2):
        a.append(x)
        b = []
        for y in gen_range(3):
            b.append(y)
            c = []
            for z in gen_range(4):
                c.append(z)
            assert c == [0, 1, 2, 3]
        assert b == [0, 1, 2]
    assert a == [0, 1]
    assert list(gen_range(5)) == list(range(5))

def gen_a(x: int) -> Iterator[int]:
    yield x + 1

def gen_b(x: int) -> Iterator[int]:
    yield x + 2

def test_generator_identities() -> None:
    # Sanity check: two distinct live objects can't reuse the same memory location
    g1 = gen_a(1)
    g2 = gen_a(1)
    assert g1 is not g2

    # If two generators have non-overlapping lifetimes, they should reuse a memory location
    g3 = gen_b(1)
    id1 = id(g3)
    g3 = gen_b(1)
    assert id(g3) == id1

    # More complex case of reuse: allocate other objects in between
    g4: Any = gen_a(1)
    id2 = id(g4)
    g4 = gen_b(1)
    g4 = [gen_b(n) for n in range(100)]
    g4 = gen_a(1)
    assert id(g4) == id2

[case testGeneratorReuseWithGilDisabled]
import sys
import threading
from typing import Iterator

def gen() -> Iterator[int]:
    yield 1

def is_gil_disabled() -> bool:
    return hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled()

def test_each_thread_gets_separate_instance() -> None:
    if not is_gil_disabled():
        # This only makes sense if GIL is disabled
        return

    g = gen()
    id1 = id(g)

    id2 = 0

    def run() -> None:
        nonlocal id2
        g = gen()
        id2 = id(g)

    t = threading.Thread(target=run)
    t.start()
    t.join()

    # Each thread should get a separate reused instance
    assert id1 != id2

[case testGeneratorWithUndefinedLocalInEnvironment]
from typing import Iterator

from testutil import assertRaises

def gen(set: bool) -> Iterator[float]:
    if set:
        y = float("-113.0")
    yield 1.0
    yield y

def test_bitmap_is_cleared_when_object_is_reused() -> None:
    # This updates the bitmap of the shared instance.
    list(gen(True))

    # Ensure bitmap has been cleared.
    with assertRaises(AttributeError):  # TODO: Should be UnboundLocalError
      list(gen(False))

def gen2(set: bool) -> Iterator[int]:
    if set:
        y = int("5")
    yield 1
    yield y

def test_undefined_int_in_environment() -> None:
    list(gen2(True))

    with assertRaises(AttributeError):  # TODO: Should be UnboundLocalError
      list(gen2(False))

[case testVariableWithSameNameAsHelperMethod]
from testutil import assertRaises
from typing import Iterator

def gen_send() -> Iterator[int]:
    send = 1
    yield send + 1

def gen_throw() -> Iterator[int]:
    throw = 42
    yield throw * 2

def undefined() -> Iterator[int]:
    if int():
        send = 1
    yield send + 1

def test_same_names() -> None:
    assert list(gen_send()) == [2]
    assert list(gen_throw()) == [84]

    with assertRaises(AttributeError, "attribute 'send' of 'undefined_gen' undefined"):
        # TODO: Should be UnboundLocalError, this test verifies that the attribute name
        # matches the variable name in the input code, since internally it's generated
        # with a prefix.
        list(undefined())

[case testGeneratorInheritance]
from typing import Iterator

class Base1:
    def foo(self) -> Iterator[int]:
        yield 1

class Derived1(Base1):
    def foo(self) -> Iterator[int]:
        yield 2
        yield 3

def base1_foo(b: Base1) -> list[int]:
    a = []
    for x in b.foo():
        a.append(x)
    return a

def derived1_foo(b: Derived1) -> list[int]:
    a = []
    for x in b.foo():
        a.append(x)
    return a

def test_generator_override() -> None:
     assert base1_foo(Base1()) == [1]
     assert base1_foo(Derived1()) == [2, 3]
     assert derived1_foo(Derived1()) == [2, 3]

[case testGeneratorEmptyTuple]
from collections.abc import Generator
from typing import Optional, Union

def test_compiledGeneratorEmptyTuple() -> None:
    jobs: Generator[Optional[str], None, None] = (_ for _ in ())
    assert list(jobs) == []