# 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

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)

[file driver.py]
from native import normal, generator, triple, another_triple, outer
from testutil import run_generator

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)

[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: bool

class C:
    foo: Foo

    def genf(self) -> Iterator[None]:
        self.foo.flag = True
        yield
        self.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

[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