b16f665a81
## Summary This PR adds support for unpacking tuple expression in an assignment statement where the target expression can be a tuple or a list (the allowed sequence targets). The implementation introduces a new `infer_assignment_target` which can then be used for other targets like the ones in for loops as well. This delegates it to the `infer_definition`. The final implementation uses a recursive function that visits the target expression in source order and compares the variable node that corresponds to the definition. At the same time, it keeps track of where it is on the assignment value type. The logic also accounts for the number of elements on both sides such that it matches even if there's a gap in between. For example, if there's a starred expression like `(a, *b, c) = (1, 2, 3)`, then the type of `a` will be `Literal[1]` and the type of `b` will be `Literal[2]`. There are a couple of follow-ups that can be done: * Use this logic for other target positions like `for` loop * Add diagnostics for mis-match length between LHS and RHS ## Test Plan Add various test cases using the new markdown test framework. Validate that existing test cases pass. --------- Co-authored-by: Carl Meyer <carl@astral.sh>
6.9 KiB
6.9 KiB
Unpacking
Tuple
Simple tuple
(a, b, c) = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
Simple list
[a, b, c] = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
Simple mixed
[a, (b, c), d] = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
reveal_type(d) # revealed: Literal[4]
Multiple assignment
a, b = c = 1, 2
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: tuple[Literal[1], Literal[2]]
Nested tuple with unpacking
(a, (b, c), d) = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Literal[3]
reveal_type(d) # revealed: Literal[4]
Nested tuple without unpacking
(a, b, c) = (1, (2, 3), 4)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: tuple[Literal[2], Literal[3]]
reveal_type(c) # revealed: Literal[4]
Uneven unpacking (1)
# TODO: Add diagnostic (there aren't enough values to unpack)
(a, b, c) = (1, 2)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
reveal_type(c) # revealed: Unknown
Uneven unpacking (2)
# TODO: Add diagnostic (too many values to unpack)
(a, b) = (1, 2, 3)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
Starred expression (1)
# TODO: Add diagnostic (need more values to unpack)
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c, d] = (1, 2) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[2]
reveal_type(d) # revealed: Unknown
Starred expression (2)
[a, *b, c] = (1, 2) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[2]
Starred expression (3)
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c] = (1, 2, 3) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[3]
Starred expression (4)
# TODO: Remove 'not-iterable' diagnostic
[a, *b, c, d] = (1, 2, 3, 4, 5, 6) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: Literal[5]
reveal_type(d) # revealed: Literal[6]
Starred expression (5)
# TODO: Remove 'not-iterable' diagnostic
[a, b, *c] = (1, 2, 3, 4) # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: Literal[2]
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(c) # revealed: @Todo
Non-iterable unpacking
TODO: Remove duplicate diagnostics. This is happening because for a sequence-like assignment target, multiple definitions are created and the inference engine runs on each of them which results in duplicate diagnostics.
# error: "Object of type `Literal[1]` is not iterable"
# error: "Object of type `Literal[1]` is not iterable"
a, b = 1
reveal_type(a) # revealed: Unknown
reveal_type(b) # revealed: Unknown
Custom iterator unpacking
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
(a, b) = Iterable()
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int
Custom iterator unpacking nested
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
(a, (b, c), d) = (1, Iterable(), 2)
reveal_type(a) # revealed: Literal[1]
reveal_type(b) # revealed: int
reveal_type(c) # revealed: int
reveal_type(d) # revealed: Literal[2]
String
Simple unpacking
a, b = 'ab'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
Uneven unpacking (1)
# TODO: Add diagnostic (there aren't enough values to unpack)
a, b, c = 'ab'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
reveal_type(c) # revealed: Unknown
Uneven unpacking (2)
# TODO: Add diagnostic (too many values to unpack)
a, b = 'abc'
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
Starred expression (1)
# TODO: Add diagnostic (need more values to unpack)
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c, d) = "ab" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
reveal_type(d) # revealed: Unknown
Starred expression (2)
(a, *b, c) = "ab" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
Starred expression (3)
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c) = "abc" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
Starred expression (4)
# TODO: Remove 'not-iterable' diagnostic
(a, *b, c, d) = "abcdef" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b) # revealed: @Todo
reveal_type(c) # revealed: LiteralString
reveal_type(d) # revealed: LiteralString
Starred expression (5)
# TODO: Remove 'not-iterable' diagnostic
(a, b, *c) = "abcd" # error: "Object of type `None` is not iterable"
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
# TODO: Should be list[int] once support for assigning to starred expression is added
reveal_type(c) # revealed: @Todo