# Test cases for "range" objects, "for" and "while" loops (compile and run)

[case testFor]
from typing import Any, List, Tuple
def count(n: int) -> None:
    for i in range(n):
        print(i)
def count_between(n: int, k: int) -> None:
    for i in range(n, k):
        print(i)
    print('n=', n)
def count_down(n: int, k: int) -> None:
    for i in range(n, k, -1):
        print(i)
def count_double(n: int, k: int) -> None:
    for i in range(n, k, 2):
        print(i)
def list_iter(l: List[int]) -> None:
    for i in l:
        print(i)
def tuple_iter(l: Tuple[int, ...]) -> None:
    for i in l:
        print(i)
def empty_tuple_iter(l: Tuple[()]) -> None:
    i: Any
    for i in l:
        print(i)
def str_iter(l: str) -> None:
    for i in l:
        print(i)
def list_rev_iter(l: List[int]) -> None:
    for i in reversed(l):
        print(i)
def list_rev_iter_lol(l: List[int]) -> None:
    for i in reversed(l):
        print(i)
        if i == 3:
            while l:
                l.pop()
def count_down_short() -> None:
    for i in range(10, 0, -1):
        print(i)
[file driver.py]
from native import (
    count, list_iter, list_rev_iter, list_rev_iter_lol, count_between, count_down, count_double,
    count_down_short, tuple_iter, empty_tuple_iter, str_iter,
)
count(5)
list_iter(list(reversed(range(5))))
list_rev_iter(list(reversed(range(5))))
count_between(11, 15)
count_between(10**20, 10**20+3)
count_down(20, 10)
count_double(10, 15)
count_down_short()
print('==')
list_rev_iter_lol(list(reversed(range(5))))
tuple_iter((1, 2, 3))
empty_tuple_iter(())
str_iter("abc")
[out]
0
1
2
3
4
4
3
2
1
0
0
1
2
3
4
11
12
13
14
n= 11
100000000000000000000
100000000000000000001
100000000000000000002
n= 100000000000000000000
20
19
18
17
16
15
14
13
12
11
10
12
14
10
9
8
7
6
5
4
3
2
1
==
0
1
2
3
1
2
3
a
b
c

[case testLoopElse]
from typing import Iterator
def run_for_range(n: int) -> None:
    for i in range(n):
        if i == 3:
            break
        print(i)
    else:
        print(n+1)

def run_for_list(n: int) -> None:
    for i in list(range(n)):
        if i == 3:
            break
        print(i)
    else:
        print(n+1)

def run_for_iter(n: int) -> None:
    def identity(x: Iterator[int]) -> Iterator[int]:
        return x
    for i in identity(range(n)):
        if i == 3:
            break
        print(i)
    else:
        print(n+1)

def count(n: int) -> int:
    i = 1
    while i <= n:
        i = i + 1
        if i == 5:
            break
    else:
        i *= -1
    return i

def nested_while() -> int:
    while True:
        while False:
            pass
        else:
            break
    else:
        return -1
    return 0

def nested_for() -> int:
    for x in range(1000):
        for y in [1,2,3]:
            pass
        else:
            break
    else:
        return -1
    return 0

[file driver.py]
from native import run_for_range, run_for_list, run_for_iter, count, nested_while, nested_for
assert nested_while() == 0
assert nested_for() == 0
assert count(0) == -1
assert count(1) == -2
assert count(5) == 5
assert count(6) == 5
run_for_range(3)
run_for_range(5)
print('==')
run_for_list(3)
run_for_list(5)
print('==')
run_for_iter(3)
run_for_iter(5)
[out]
0
1
2
4
0
1
2
==
0
1
2
4
0
1
2
==
0
1
2
4
0
1
2

[case testNestedLoopSameIdx]
from typing import List, Generator

def nested_enumerate() -> None:
    l1 = [0,1,2]
    l2 = [0,1,2]
    outer_seen = []
    outer = 0
    for i, j in enumerate(l1):
        assert i == outer
        outer_seen.append(i)
        inner = 0
        for i, k in enumerate(l2):
            assert i == inner
            inner += 1
        outer += 1
    assert i == 2
    assert outer_seen == l1

def nested_range() -> None:
    outer = 0
    outer_seen = []
    for i in range(3):
        assert i == outer
        outer_seen.append(i)
        inner = 0
        for i in range(3):
            assert i == inner
            inner += 1
        outer += 1
    assert outer_seen == [0,1,2]

def nested_list() -> None:
    l1 = [0,1,2]
    l2 = [0,1,2]
    outer_seen = []
    outer = 0
    for i in l1:
        assert i == outer
        outer_seen.append(i)
        inner = 0
        for i in l2:
            assert i == inner
            inner += 1
        outer += 1
    assert outer_seen == l1

def nested_yield() -> Generator:
    for i in range(3):
        for i in range(3):
            yield i
        yield i


[file driver.py]
from native import nested_enumerate, nested_range, nested_list, nested_yield
nested_enumerate()
nested_range()
nested_list()
gen = nested_yield()
expected = [0, 1, 2, 2, 0, 1, 2, 2, 0, 1, 2, 2]
assert list(gen) == expected
[out]

[case testRangeLoopVariableAfterLoop]
for i in range(3):
    pass
assert i == 2, f"expected 2, got {i}"

for j in range(0, 10, 3):
    pass
assert j == 9, f"expected 9, got {j}"

for k in range(5, 0, -1):
    pass
assert k == 1, f"expected 1, got {k}"

for m in range(0, 10, 2):
    pass
assert m == 8, f"expected 8, got {m}"

[file driver.py]
from native import *
[out]

[case testForIterable]
from typing import Iterable, Dict, Any, Tuple, TypeVar

T = TypeVar("T")

def iterate_over_any(a: Any) -> None:
    for element in a:
        print(element)

def iterate_over_iterable(iterable: Iterable[T]) -> None:
    for element in iterable:
        print(element)

def iterate_and_delete(d: Dict[int, int]) -> None:
    for key in d:
        d.pop(key)

def sum_over_values(d: Dict[int, int]) -> int:
    s = 0
    for key in d:
        s = s + d[key]
    return s

def sum_over_even_values(d: Dict[int, int]) -> int:
    s = 0
    for key in d:
        if d[key] % 2:
            continue
        s = s + d[key]
    return s

def sum_over_two_values(d: Dict[int, int]) -> int:
    s = 0
    i = 0
    for key in d:
        if i == 2:
            break
        s = s + d[key]
        i = i + 1
    return s

def iterate_over_tuple(iterable: Tuple[int, int, int]) -> None:
    for element in iterable:
        print(element)

[file driver.py]
from native import iterate_over_any, iterate_over_iterable, iterate_and_delete, sum_over_values, sum_over_even_values, sum_over_two_values, iterate_over_tuple
import traceback
def broken_generator(n):
    num = 0
    while num < n:
        yield num
        num += 1
    raise Exception('Exception Manually Raised')

d = {1:1, 2:2, 3:3, 4:4, 5:5}
print(sum_over_values(d))
print(sum_over_even_values(d))
print(sum_over_two_values(d))

try:
    iterate_over_any(5)
except TypeError:
    traceback.print_exc()
try:
    iterate_over_iterable(broken_generator(5))
except Exception:
    traceback.print_exc()
try:
    iterate_and_delete(d)
except RuntimeError:
    traceback.print_exc()

iterate_over_tuple((1, 2, 3))
[out]
Traceback (most recent call last):
  File "driver.py", line 16, in <module>
    iterate_over_any(5)
  File "native.py", line 6, in iterate_over_any
    for element in a:
TypeError: 'int' object is not iterable
Traceback (most recent call last):
  File "driver.py", line 20, in <module>
    iterate_over_iterable(broken_generator(5))
  File "native.py", line 10, in iterate_over_iterable
    for element in iterable:
  File "driver.py", line 8, in broken_generator
    raise Exception('Exception Manually Raised')
Exception: Exception Manually Raised
Traceback (most recent call last):
  File "driver.py", line 24, in <module>
    iterate_and_delete(d)
  File "native.py", line 14, in iterate_and_delete
    for key in d:
RuntimeError: dictionary changed size during iteration
15
6
3
0
1
2
3
4
1
2
3
[out version>=3.13]
Traceback (most recent call last):
  File "driver.py", line 16, in <module>
    iterate_over_any(5)
    ~~~~~~~~~~~~~~~~^^^
  File "native.py", line 6, in iterate_over_any
    for element in a:
TypeError: 'int' object is not iterable
Traceback (most recent call last):
  File "driver.py", line 20, in <module>
    iterate_over_iterable(broken_generator(5))
    ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
  File "native.py", line 10, in iterate_over_iterable
    for element in iterable:
  File "driver.py", line 8, in broken_generator
    raise Exception('Exception Manually Raised')
Exception: Exception Manually Raised
Traceback (most recent call last):
  File "driver.py", line 24, in <module>
    iterate_and_delete(d)
    ~~~~~~~~~~~~~~~~~~^^^
  File "native.py", line 14, in iterate_and_delete
    for key in d:
RuntimeError: dictionary changed size during iteration
15
6
3
0
1
2
3
4
1
2
3

[case testContinueFor]
def f() -> None:
    for n in range(5):
        continue
[file driver.py]
from native import f
f()

[case testMultipleVarsWithLoops]
# Test comprehensions and for loops with multiple index variables
l = [(1, 2, 'a'), (3, 4, 'b'), (5, 6, 'c')]
l2 = [str(a*100+b)+' '+c for a, b, c in l]
l3 = []
for a, b, c in l:
    l3.append(str(a*1000+b)+' '+c)
[file driver.py]
from native import l, l2, l3
for a in l2 + l3:
    print(a)
[out]
102 a
304 b
506 c
1002 a
3004 b
5006 c

[case testForZipAndEnumerate]
from typing import Iterable, List, Any
def f(a: Iterable[int], b: List[int]) -> List[Any]:
    res = []
    for (x, y), z in zip(enumerate(a), b):
        res.append((x, y, z))
    return res
def g(a: Iterable[int], b: Iterable[str]) -> List[Any]:
    res = []
    for x, (y, z) in enumerate(zip(a, b)):
        res.append((x, y, z))
    return res

[file driver.py]
from native import f, g

assert f([6, 7], [8, 9]) == [(0, 6, 8), (1, 7, 9)]
assert g([6, 7], ['a', 'b']) == [(0, 6, 'a'), (1, 7, 'b')]
assert f([6, 7], [8]) == [(0, 6, 8)]
assert f([6], [8, 9]) == [(0, 6, 8)]

[case testEnumerateEmptyList]
from typing import List

def get_enumerate_locals(iterable: List[int]) -> int:
    for i, j in enumerate(iterable):
        pass
    try:
        return i
    except NameError:
        return -100

[file driver.py]
from native import get_enumerate_locals

print(get_enumerate_locals([]))
print(get_enumerate_locals([55]))
print(get_enumerate_locals([551, 552]))

[out]
-100
0
1

[case testIterTypeTrickiness]
# Test inferring the type of a for loop body doesn't cause us grief
# Extracted from somethings that broke in mypy

from typing import Optional

# really I only care that this one build
def foo(x: object) -> None:
    if isinstance(x, dict):
        for a in x:
            pass

def bar(x: Optional[str]) -> None:
    vars = (
        ("a", 'lol'),
        ("b", 'asdf'),
        ("lol", x),
        ("an int", 10),
    )
    for name, value in vars:
        pass

[file driver.py]
from native import bar
bar(None)

[case testRangeObject]
from typing import Any

def f(x: range) -> int:
    sum = 0
    for i in x:
        sum += i
    return sum

def test_range_object() -> None:
    r1 = range(4, 12, 2)
    tmp_list = [x for x in r1]
    assert tmp_list == [4, 6, 8, 10]
    assert f(r1) == 28
    r2: Any = range(10)
    assert f(r2) == 45
    r3: Any = 'x'
    try:
        f(r3)
    except TypeError as e:
        assert "range object expected; got str" in str(e)
    try:
        ff: Any = f
        ff(r3)
    except TypeError as e:
        assert "range object expected; got str" in str(e)
    try:
        r4 = range(4, 12, 0)
    except ValueError as e:
        assert "range() arg 3 must not be zero" in str(e)

[case testNamedTupleLoop]
from collections.abc import Iterable
from typing import NamedTuple, Any
from typing_extensions import Self


class Vector2(NamedTuple):
    x: int
    y: float

    @classmethod
    def from_iter(cls, iterable: Iterable[Any]) -> Self:
        return cls(*iter(iterable))

    def __neg__(self) -> Self:
        return self.from_iter(-c for c in self)

[file driver.py]
import native
print(-native.Vector2(2, -3.1))
print([x for x in native.Vector2(4, -5.2)])

[out]
Vector2(x=-2, y=3.1)
\[4, -5.2]