[red-knot] Add support for unpacking for target (#15058)

## Summary

Related to #13773 

This PR adds support for unpacking `for` statement targets.

This involves updating the `value` field in the `Unpack` target to use
an enum which specifies the "where did the value expression came from?".
This is because for an iterable expression, we need to unpack the
iterator type while for assignment statement we need to unpack the value
type itself. And, this needs to be done in the unpack query.

### Question

One of the ways unpacking works in `for` statement is by looking at the
union of the types because if the iterable expression is a tuple then
the iterator type will be union of all the types in the tuple. This
means that the test cases that will test the unpacking in `for`
statement will also implicitly test the unpacking union logic. I was
wondering if it makes sense to merge these cases and only add the ones
that are specific to the union unpacking or for statement unpacking
logic.

## Test Plan

Add test cases involving iterating over a tuple type. I've intentionally
left out certain cases for now and I'm curious to know any thoughts on
the above query.
This commit is contained in:
Dhruv Manilawala
2024-12-23 11:43:49 +05:30
committed by GitHub
parent b6c8f5d79e
commit 113c804a62
7 changed files with 294 additions and 83 deletions

View File

@@ -472,3 +472,104 @@ def _(arg: tuple[int, str] | Iterable):
reveal_type(a) # revealed: int | bytes
reveal_type(b) # revealed: str | bytes
```
## For statement
Unpacking in a `for` statement.
### Same types
```py
def _(arg: tuple[tuple[int, int], tuple[int, int]]):
for a, b in arg:
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int
```
### Mixed types (1)
```py
def _(arg: tuple[tuple[int, int], tuple[int, str]]):
for a, b in arg:
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int | str
```
### Mixed types (2)
```py
def _(arg: tuple[tuple[int, str], tuple[str, int]]):
for a, b in arg:
reveal_type(a) # revealed: int | str
reveal_type(b) # revealed: str | int
```
### Mixed types (3)
```py
def _(arg: tuple[tuple[int, int, int], tuple[int, str, bytes], tuple[int, int, str]]):
for a, b, c in arg:
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int | str
reveal_type(c) # revealed: int | bytes | str
```
### Same literal values
```py
for a, b in ((1, 2), (3, 4)):
reveal_type(a) # revealed: Literal[1, 3]
reveal_type(b) # revealed: Literal[2, 4]
```
### Mixed literal values (1)
```py
for a, b in ((1, 2), ("a", "b")):
reveal_type(a) # revealed: Literal[1] | Literal["a"]
reveal_type(b) # revealed: Literal[2] | Literal["b"]
```
### Mixed literals values (2)
```py
# error: "Object of type `Literal[1]` is not iterable"
# error: "Object of type `Literal[2]` is not iterable"
# error: "Object of type `Literal[4]` is not iterable"
for a, b in (1, 2, (3, "a"), 4, (5, "b"), "c"):
reveal_type(a) # revealed: Unknown | Literal[3, 5] | LiteralString
reveal_type(b) # revealed: Unknown | Literal["a", "b"]
```
### Custom iterator (1)
```py
class Iterator:
def __next__(self) -> tuple[int, int]:
return (1, 2)
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
for a, b in Iterable():
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int
```
### Custom iterator (2)
```py
class Iterator:
def __next__(self) -> bytes:
return b""
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
def _(arg: tuple[tuple[int, str], Iterable]):
for a, b in arg:
reveal_type(a) # revealed: int | bytes
reveal_type(b) # revealed: str | bytes
```