ruff/crates/red_knot_python_semantic/resources/mdtest/comprehensions/basic.md

Comprehensions

Basic comprehensions

class IntIterator:
    def __next__(self) -> int:
        return 42

class IntIterable:
    def __iter__(self) -> IntIterator:
        return IntIterator()

# revealed: int
[reveal_type(x) for x in IntIterable()]

class IteratorOfIterables:
    def __next__(self) -> IntIterable:
        return IntIterable()

class IterableOfIterables:
    def __iter__(self) -> IteratorOfIterables:
        return IteratorOfIterables()

# revealed: tuple[int, IntIterable]
[reveal_type((x, y)) for y in IterableOfIterables() for x in y]

# revealed: int
{reveal_type(x): 0 for x in IntIterable()}

# revealed: int
{0: reveal_type(x) for x in IntIterable()}

Nested comprehension

class IntIterator:
    def __next__(self) -> int:
        return 42

class IntIterable:
    def __iter__(self) -> IntIterator:
        return IntIterator()

# TODO: This could be a `tuple[int, int]` if we model that `y` can not be modified in the outer comprehension scope
# revealed: tuple[int, Unknown | int]
[[reveal_type((x, y)) for x in IntIterable()] for y in IntIterable()]

Comprehension referencing outer comprehension

class IntIterator:
    def __next__(self) -> int:
        return 42

class IntIterable:
    def __iter__(self) -> IntIterator:
        return IntIterator()

class IteratorOfIterables:
    def __next__(self) -> IntIterable:
        return IntIterable()

class IterableOfIterables:
    def __iter__(self) -> IteratorOfIterables:
        return IteratorOfIterables()

# TODO: This could be a `tuple[int, int]` (see above)
# revealed: tuple[int, Unknown | IntIterable]
[[reveal_type((x, y)) for x in y] for y in IterableOfIterables()]

Comprehension with unbound iterable

Iterating over an unbound iterable yields Unknown:

# error: [unresolved-reference] "Name `x` used when not defined"
# revealed: Unknown
[reveal_type(z) for z in x]

class IntIterator:
    def __next__(self) -> int:
        return 42

class IntIterable:
    def __iter__(self) -> IntIterator:
        return IntIterator()

# error: [not-iterable] "Object of type `int` is not iterable"
# revealed: tuple[int, Unknown]
[reveal_type((x, z)) for x in IntIterable() for z in x]

Starred expressions

Starred expressions must be iterable

class NotIterable: ...

class Iterator:
    def __next__(self) -> int:
        return 42

class Iterable:
    def __iter__(self) -> Iterator: ...

# This is fine:
x = [*Iterable()]

# error: [not-iterable] "Object of type `NotIterable` is not iterable"
y = [*NotIterable()]

Async comprehensions

Basic

class AsyncIterator:
    async def __anext__(self) -> int:
        return 42

class AsyncIterable:
    def __aiter__(self) -> AsyncIterator:
        return AsyncIterator()

# revealed: @Todo(async iterables/iterators)
[reveal_type(x) async for x in AsyncIterable()]

Invalid async comprehension

This tests that we understand that async comprehensions do not work according to the synchronous iteration protocol

class Iterator:
    def __next__(self) -> int:
        return 42

class Iterable:
    def __iter__(self) -> Iterator:
        return Iterator()

# revealed: @Todo(async iterables/iterators)
[reveal_type(x) async for x in Iterable()]