crates/ty_python_semantic/resources/mdtest/comprehensions/basic.md
# revealed: int
[reveal_type(x) for x in range(3)]
class Row:
def __next__(self) -> range:
return range(3)
class Table:
def __iter__(self) -> Row:
return Row()
# revealed: tuple[int, range]
[reveal_type((cell, row)) for row in Table() for cell in row]
# revealed: int
{reveal_type(x): 0 for x in range(3)}
# revealed: int
{0: reveal_type(x) for x in range(3)}
# revealed: tuple[int, int]
[[reveal_type((x, y)) for x in range(3)] for y in range(3)]
PEP 572 specifies that an assignment expression in a comprehension binds its target in the scope containing the outermost comprehension.
ty currently assumes that a comprehension runs at least once and that a generator expression is consumed immediately.
Assignment expressions can appear in the element of a list comprehension and in the key or value of a dictionary comprehension:
[(list_value := item) for item in [1]]
{(dict_key := item): (dict_value := item) for item in [1]}
reveal_type(list_value) # revealed: int
reveal_type(dict_key) # revealed: int
reveal_type(dict_value) # revealed: int
The target also binds in the containing scope when the assignment is in a generator expression. PEP
572 uses this any pattern as a motivating example:
def find_comment(lines: list[str]):
if any((comment := line).startswith("#") for line in lines):
reveal_type(comment) # revealed: str
If an iteration assigns the same target more than once, the last assignment determines its value after the comprehension:
[(ordered := item, ordered := "") for item in [1]]
reveal_type(ordered) # revealed: str
A target in a branch known not to run remains unbound, while the other target is available after the comprehension:
[(dead := 1) if False else (live := 2) for _ in [0]]
dead # error: [unresolved-reference]
reveal_type(live) # revealed: int
When the assignment only runs on one possible path, an earlier value remains possible:
def conditional_with_previous_value(flag: bool):
value = "old"
[(value := 1) if flag else 0 for _ in [0]]
reveal_type(value) # revealed: Literal["old"] | int
Without an earlier value, the target may be unbound:
def conditional_without_previous_value(flag: bool):
[(value := 1) if flag else 0 for _ in [0]]
# error: [possibly-unresolved-reference]
reveal_type(value) # revealed: int
ty conservatively keeps the type of an assignment that is unreachable on the first iteration, since
a later iteration may take a different branch. Even though 0 == 1 is always false, the target is
therefore possibly unbound, and checking must continue after the read:
def statically_false_condition():
[(value := 1) if 0 == 1 else 0 for _ in [0]]
# error: [possibly-unresolved-reference]
reveal_type(value) # revealed: int
still_reachable # error: [unresolved-reference]
A false filter skips the element, but an assignment made while evaluating that filter still takes effect:
[value for value in [True, False] if (last_value := value)]
reveal_type(last_value) # revealed: bool
If short-circuit evaluation skips the assignment, the target may be unbound:
def conditional_filter(flag: bool):
[0 for _ in [0] if flag and (value := 1)]
# error: [possibly-unresolved-reference]
reveal_type(value) # revealed: int
An assignment in the element only runs when every preceding filter succeeds:
def assignment_after_filter(flag: bool):
[(value := 1) for _ in [0] if flag]
# error: [possibly-unresolved-reference]
reveal_type(value) # revealed: int
An assignment can read the value left by an earlier iteration. In this example, the final value is
3, so retaining only the first iteration's literal values would be incorrect:
def partial_sum():
total = 0
[total := total + value for value in [1, 2]]
reveal_type(total) # revealed: int
ty does not yet account for a type that changes between iterations. The second iteration below
assigns int, so the final type should be str | int and value.upper() should report an error:
def type_changes_across_iterations():
value = 0
[value := "" if isinstance(value, int) else 0 for _ in [0, 1]]
reveal_type(value) # revealed: str
value.upper()
The same applies when two targets depend on values from earlier iterations:
def two_dependent_targets():
x = 0
y = 0
[(y := x, x := y + 1) for _ in [1, 2]]
reveal_type(x) # revealed: int
reveal_type(y) # revealed: int
A guard can also depend on a value changed by a later assignment in the same iteration. The first
iteration below sets flag, so the second iteration assigns value:
def loop_carried_guard():
flag = False
[((value := 1) if flag else 0, (flag := True)) for _ in [0, 1]]
# error: [possibly-unresolved-reference]
reveal_type(value) # revealed: int
An assignment in a branch known not to run still makes its target local to the containing function. A read must not fall back to a global variable with the same name:
local_target = "global"
def read_local_target():
[(local_target := 1) if False else 0 for _ in [0]]
local_target # error: [unresolved-reference]
A walrus also makes its target local before the first iteration. Its first assignment must not read
a global with the same name. Explicit global and nonlocal declarations still refer to the
existing outer variable:
total = 0
def sums(values: list[int]) -> list[int]:
return [total := total + value for value in values] # error: [unresolved-reference]
def sums_global(values: list[int]) -> list[int]:
global total
return [total := total + value for value in values]
def sums_nonlocal(values: list[int]) -> list[int]:
total = 0
def add_values() -> list[int]:
nonlocal total
return [total := total + value for value in values]
return add_values()
An assignment in an inner comprehension still binds outside the outermost comprehension. A later assignment in the outer comprehension replaces the inner value:
[([nested_order := 1 for _ in [0]], (nested_order := "")) for _ in [0]]
reveal_type(nested_order) # revealed: str
These are controls for an inner comprehension that is never evaluated. It must not replace an earlier value:
def unreachable_nested_assignment_with_previous_value():
value = "old"
[[value := 1 for _ in [0]] if False else [] for _ in [0]]
reveal_type(value) # revealed: Literal["old"]
Nor should it create a new value:
def unreachable_nested_assignment_without_previous_value():
[[value := 1 for _ in [0]] if False else [] for _ in [0]]
value # error: [unresolved-reference]
class Row:
def __next__(self) -> range:
return range(3)
class Table:
def __iter__(self) -> Row:
return Row()
# revealed: tuple[int, range]
[[reveal_type((cell, row)) for cell in row] for row in Table()]
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]
# error: [not-iterable] "Object of type `int` is not iterable"
# revealed: tuple[int, Unknown]
[reveal_type((x, z)) for x in range(3) for z in x]
# error: [unresolved-reference] "Name `foo` used when not defined"
foo
foo = [
# revealed: tuple[int, Unknown]
reveal_type((x, z))
for x in range(3)
# error: [unresolved-reference] "Name `foo` used when not defined"
for z in [foo]
]
baz = [
# revealed: tuple[int, Unknown]
reveal_type((x, z))
for x in range(3)
# error: [unresolved-reference] "Name `baz` used when not defined"
for z in [baz]
]
Starred expressions must be iterable
class NotIterable: ...
# This is fine:
x = [*range(3)]
# error: [not-iterable] "Object of type `NotIterable` is not iterable"
y = [*NotIterable()]
class AsyncIterator:
async def __anext__(self) -> int:
return 42
class AsyncIterable:
def __aiter__(self) -> AsyncIterator:
return AsyncIterator()
async def _():
# revealed: int
[reveal_type(x) async for x in AsyncIterable()]
This tests that we understand that async comprehensions do not work according to the synchronous
iteration protocol
async def _():
# error: [not-iterable] "Object of type `range` is not async-iterable"
# revealed: Unknown
[reveal_type(x) async for x in range(3)]
The type of the expression being iterated over is immutable, and so should not be widened with
Unknown or through literal promotion:
x = [
reveal_type(string) # revealed: Literal["a", "b"]
for string in ["a", "b"]
]
The type of the comprehension expression itself should reflect the inferred element type:
from typing import TypedDict, Literal
# revealed: list[int]
reveal_type([x for x in range(10)])
# revealed: set[int]
reveal_type({x for x in range(10)})
# revealed: dict[int, str]
reveal_type({x: str(x) for x in range(10)})
# revealed: list[tuple[int, str]]
reveal_type([(x, y) for x in range(5) for y in ["a", "b", "c"]])
squares: list[int | None] = [x**2 for x in range(10)]
reveal_type(squares) # revealed: list[int | None]
[environment]
python-version = "3.15"
Unpacking comprehensions flatten the unpacked element type:
list_of_lists: list[list[int]] = [[1], [2, 3]]
sets: list[set[str]] = [{"a"}, {"b", "c"}]
dicts: list[dict[str, int]] = [{"a": 1}, {"b": 2}]
not_iterables: list[int] = [1, 2]
reveal_type([*xs for xs in list_of_lists]) # revealed: list[int]
reveal_type({*xs for xs in sets}) # revealed: set[str]
reveal_type({**d for d in dicts}) # revealed: dict[str, int]
[*value for value in not_iterables] # error: [not-iterable] "Object of type `int` is not iterable"
{*value for value in not_iterables} # error: [not-iterable] "Object of type `int` is not iterable"
{**value for value in not_iterables} # error: [invalid-argument-type]
Inference for comprehensions takes the type context into account:
from typing import Literal, Sequence, TypedDict
# Without type context:
reveal_type([x for x in [1, 2, 3]]) # revealed: list[int]
reveal_type({x: "a" for x in [1, 2, 3]}) # revealed: dict[int, str]
reveal_type({str(x): x for x in [1, 2, 3]}) # revealed: dict[str, int]
reveal_type({x for x in [1, 2, 3]}) # revealed: set[int]
# With type context:
x1: list[int] = [x for x in [1, 2, 3]]
reveal_type(x1) # revealed: list[int]
x2: Sequence[int] = [x for x in [1, 2, 3]]
reveal_type(x2) # revealed: list[int]
x3: dict[int, str] = {x: str(x) for x in [1, 2, 3]}
reveal_type(x3) # revealed: dict[int, str]
x4: set[int] = {x for x in [1, 2, 3]}
reveal_type(x4) # revealed: set[int]
This also works for nested comprehensions:
table = [[(x, y) for x in range(3)] for y in range(3)]
reveal_type(table) # revealed: list[list[tuple[int, int]]]
table_with_content: list[list[tuple[int, int, str | None]]] = [[(x, y, None) for x in range(3)] for y in range(3)]
reveal_type(table_with_content) # revealed: list[list[tuple[int, int, str | None]]]
The type context is propagated down into the comprehension:
y1: list[list[int]] = [[n] for n in [1, 2, 3]]
reveal_type(y1) # revealed: list[list[int]]
y2: list[Sequence[int]] = [[i] for i in [1, 2, 3]]
reveal_type(y2) # revealed: list[Sequence[int]]
class Person(TypedDict):
name: str
y3: list[Person] = [{"name": n} for n in ["Alice", "Bob"]]
reveal_type(y3) # revealed: list[Person]
# error: [invalid-assignment]
# error: [invalid-key] "Unknown key "misspelled" for TypedDict `Person`"
# error: [missing-typed-dict-key] "Missing required key 'name' in TypedDict `Person` constructor"
y4: list[Person] = [{"misspelled": n} for n in ["Alice", "Bob"]]
We promote literals to avoid overly-precise types in invariant positions:
reveal_type([x for x in ("a", "b", "c")]) # revealed: list[str]
reveal_type({x for x in (1, 2, 3)}) # revealed: set[int]
reveal_type({k: 0 for k in ("a", "b", "c")}) # revealed: dict[str, int]
Type context can prevent this promotion from happening:
list_of_literals: list[Literal["a", "b", "c"]] = [x for x in ("a", "b", "c")]
reveal_type(list_of_literals) # revealed: list[Literal["a", "b", "c"]]
dict_with_literal_keys: dict[Literal["a", "b", "c"], int] = {k: 0 for k in ("a", "b", "c")}
reveal_type(dict_with_literal_keys) # revealed: dict[Literal["a", "b", "c"], int]
dict_with_literal_values: dict[str, Literal[1, 2, 3]] = {str(k): k for k in (1, 2, 3)}
reveal_type(dict_with_literal_values) # revealed: dict[str, Literal[1, 2, 3]]
set_with_literals: set[Literal[1, 2, 3]] = {k for k in (1, 2, 3)}
reveal_type(set_with_literals) # revealed: set[Literal[1, 2, 3]]