-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[red-knot] Infer target types for unpacked tuple assignment (#13316)
## 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 <[email protected]>
- Loading branch information
1 parent
d774807
commit b16f665
Showing
7 changed files
with
525 additions
and
46 deletions.
There are no files selected for viewing
273 changes: 273 additions & 0 deletions
273
crates/red_knot_python_semantic/resources/mdtest/unpacking.md
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,273 @@ | ||
# Unpacking | ||
|
||
## Tuple | ||
|
||
### Simple tuple | ||
|
||
```py | ||
(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 | ||
|
||
```py | ||
[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 | ||
|
||
```py | ||
[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 | ||
|
||
```py | ||
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 | ||
|
||
```py | ||
(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 | ||
|
||
```py | ||
(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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
[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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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. | ||
|
||
```py | ||
# 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 | ||
|
||
```py | ||
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 | ||
|
||
```py | ||
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 | ||
|
||
```py | ||
a, b = 'ab' | ||
reveal_type(a) # revealed: LiteralString | ||
reveal_type(b) # revealed: LiteralString | ||
``` | ||
|
||
### Uneven unpacking (1) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# TODO: Add diagnostic (too many values to unpack) | ||
a, b = 'abc' | ||
reveal_type(a) # revealed: LiteralString | ||
reveal_type(b) # revealed: LiteralString | ||
``` | ||
|
||
### Starred expression (1) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
(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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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) | ||
|
||
```py | ||
# 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 | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.