mypyc/test-data/run-loops.test

# 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 testRangeLoopVariableEmptyRange]
from typing import Iterator
from mypy_extensions import i64
from testutil import assertRaises

def last_range(start: int, end: int) -> int:
    i = -1
    for i in range(start, end):
        pass
    return i

def unbound_range(n: int) -> int:
    for i in range(n):
        pass
    return i

def last_range_negative_step() -> int:
    i = -1
    for i in range(0, 5, -1):
        pass
    return i

def unbound_range_negative_step() -> int:
    for i in range(0, 5, -1):
        pass
    return i

def last_range_i64(n: i64) -> i64:
    i: i64 = -1
    for i in range(n):
        pass
    return i

def unbound_range_i64(n: i64) -> i64:
    for i in range(n):
        pass
    return i

def test_empty_range_keeps_old_value() -> None:
    assert last_range(5, 0) == -1
    assert last_range(0, 3) == 2
    assert last_range_negative_step() == -1
    assert last_range_i64(0) == -1
    assert last_range_i64(3) == 2

def test_empty_range_leaves_variable_unbound() -> None:
    assert unbound_range(3) == 2
    with assertRaises(UnboundLocalError):
        unbound_range(0)
    with assertRaises(UnboundLocalError):
        unbound_range_negative_step()
    assert unbound_range_i64(3) == 2
    with assertRaises(UnboundLocalError):
        unbound_range_i64(0)

def test_empty_range_in_enumerate_and_zip() -> None:
    i = x = -1
    for i, x in enumerate(range(0)):
        pass
    assert (i, x) == (-1, -1)
    a = b = -1
    for a, b in zip(range(3), range(0)):
        pass
    assert (a, b) == (-1, -1)

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

def test_empty_range_in_generator() -> None:
    assert list(gen_range(0)) == [-1]
    assert list(gen_range(2)) == [0, 1, 1]

for module_var in range(0):
    pass

def test_empty_range_at_module_level() -> None:
    assert "module_var" not in globals()

class A:
    def __init__(self, n: int) -> None:
        for self.x in range(n):
            pass

def test_empty_range_with_attribute_target() -> None:
    assert A(2).x == 1
    with assertRaises(AttributeError):
        A(0).x

[case testForLoopTargetEvaluatedEachIteration]
class C:
    def __init__(self) -> None:
        self.assigned: list[int] = []

    @property
    def x(self) -> int:
        return -1

    @x.setter
    def x(self, value: int) -> None:
        self.assigned.append(value)

def test_attribute_target() -> None:
    c = C()
    for c.x in range(0):
        pass
    assert c.assigned == []
    for c.x in range(3):
        pass
    assert c.assigned == [0, 1, 2]

def index(log: list[int]) -> int:
    log.append(0)
    return 0

def test_index_target_in_range_loop() -> None:
    a = [-1]
    log: list[int] = []
    for a[index(log)] in range(0):
        pass
    assert a == [-1]
    assert log == []
    for a[index(log)] in range(3):
        pass
    assert a == [2]
    assert len(log) == 3

def test_index_target_in_enumerate_loop() -> None:
    a = [-1]
    log: list[int] = []
    empty: list[str] = []
    for a[index(log)], s in enumerate(empty):
        pass
    assert a == [-1]
    assert log == []
    for a[index(log)], s in enumerate("abc"):
        pass
    assert a == [2]
    assert len(log) == 3

[case testForZipAndEnumerateListClearedBeforeBody]
def index(items: list[int], calls: list[int]) -> int:
    calls.append(1)
    if len(calls) == 2:
        items.clear()
    return 0

def test_enumerate_index_target() -> None:
    items = [10, 20]
    target = [-1]
    calls: list[int] = []
    result: list[int] = []
    for target[index(items, calls)], value in enumerate(items):
        result.append(value)
    assert items == []
    assert result == [10, 20]

def test_zip_range_target() -> None:
    items = [10, 20]
    target = [-1]
    calls: list[int] = []
    result: list[int] = []
    for target[index(items, calls)], value in zip(range(2), items):
        result.append(value)
    assert items == []
    assert result == [10, 20]

def test_zip_reversed() -> None:
    items = [10, 20]
    target = [-1]
    calls: list[int] = []
    result: list[int] = []
    for target[index(items, calls)], value in zip([1, 2], reversed(items)):
        result.append(value)
    assert items == []
    assert result == [20, 10]

def test_zip_iterator_clears_list() -> None:
    items = [10, 20]

    def clear_items(n: int) -> int:
        if n == 2:
            items.clear()
        return n

    result: list[tuple[int, int]] = []
    for value, n in zip(items, map(clear_items, [1, 2, 3])):
        result.append((value, n))
    assert items == []
    assert result == [(10, 1), (20, 2)]

class C:
    def __init__(self, items: list[int]) -> None:
        self.items = items

    @property
    def x(self) -> int:
        return -1

    @x.setter
    def x(self, value: int) -> None:
        if value == 1:
            self.items.clear()

def test_enumerate_index_setter_clears_list() -> None:
    items = [10, 20]
    c = C(items)
    result: list[int] = []
    for c.x, value in enumerate(items):
        result.append(value)
    assert items == []
    assert result == [10, 20]

[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 testRangeLoopIndexArgs]
from typing import Any

from testutil import assertRaises

def loop_end(end: Any) -> list[int]:
    result = []
    for i in range(end):
        result.append(i)
    return result

def loop(start: Any, end: Any) -> list[int]:
    return [i for i in range(start, end)]

def loop_down(start: Any) -> list[int]:
    return [i for i in range(start, 0, -1)]

def loop_down_to(start: Any, end: Any) -> list[int]:
    return [i for i in range(start, end, -1)]

def loop_bool(b: bool) -> list[int]:
    return [i for i in range(b)]

class NativeIndex:
    def __init__(self, n: int) -> None:
        self.n = n

    def __index__(self) -> int:
        return self.n

class NativeIndexSub(NativeIndex):
    def __index__(self) -> int:
        return self.n + 1

def loop_native(start: NativeIndex, end: NativeIndex) -> list[int]:
    # The test fixtures only accept int arguments for range()
    return [i for i in range(start, end)]  # type: ignore[arg-type]

def test_index_args() -> None:
    from index_types import Index
    assert loop_end(Index(3)) == [0, 1, 2]
    assert loop(Index(1), Index(3)) == [1, 2]
    assert loop_down(Index(2)) == [2, 1]

def test_big_index_args() -> None:
    from index_types import Index
    # Ints >= 2**62 or < -2**62 don't fit in a short int, so these loops cross
    # between short and heap-allocated values
    n = 2**62
    assert loop(Index(n - 1), Index(n + 2)) == [n - 1, n, n + 1]
    assert loop(Index(-n - 2), Index(-n + 1)) == [-n - 2, -n - 1, -n]
    assert loop_down_to(Index(n + 1), Index(n - 2)) == [n + 1, n, n - 1]
    assert loop_down_to(Index(-n + 1), Index(-n - 2)) == [-n + 1, -n, -n - 1]
    assert loop_end(Index(-n - 1)) == []

def test_native_index_args() -> None:
    assert loop_native(NativeIndex(1), NativeIndex(3)) == [1, 2]
    assert loop_native(NativeIndexSub(1), NativeIndex(4)) == [2, 3]
    n = 2**62
    assert loop_native(NativeIndex(n - 1), NativeIndex(n + 1)) == [n - 1, n]

def test_int_args() -> None:
    assert loop_end(3) == [0, 1, 2]
    assert loop(True, 3) == [1, 2]
    n = 2**100
    assert loop(n, n + 2) == [n, n + 1]
    assert loop_bool(True) == [0]
    assert loop_bool(False) == []

def test_invalid_args() -> None:
    with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"):
        loop_end(2.5)
    with assertRaises(TypeError, "'str' object cannot be interpreted as an integer"):
        loop(0, "3")
    with assertRaises(TypeError, "'float' object cannot be interpreted as an integer"):
        loop_down(2.0)

[file index_types.py]
class Index:
    def __init__(self, n: int) -> None:
        self.n = n

    def __index__(self) -> int:
        return self.n

[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]