[red-knot] Add support for unpacking union types (#15052)

## Summary

Refer:
https://github.com/astral-sh/ruff/issues/13773#issuecomment-2548020368

This PR adds support for unpacking union types. 

Unpacking a union type requires us to first distribute the types for all
the targets that are involved in an unpacking. For example, if there are
two targets and a union type that needs to be unpacked, each target will
get a type from each element in the union type.

For example, if the type is `tuple[int, int] | tuple[int, str]` and the
target has two elements `(a, b)`, then
* The type of `a` will be a union of `int` and `int` which are at index
0 in the first and second tuple respectively which resolves to an `int`.
* Similarly, the type of `b` will be a union of `int` and `str` which
are at index 1 in the first and second tuple respectively which will be
`int | str`.

### Refactors

There are couple of refactors that are added in this PR:
* Add a `debug_assertion` to validate that the unpack target is a list
or a tuple
* Add a separate method to handle starred expression

## Test Plan

Update `unpacking.md` with additional test cases that uses union types.
This is done using parameter type hints style.
This commit is contained in:
Dhruv Manilawala
2024-12-20 16:31:15 +05:30
committed by GitHub
parent 089a98e904
commit d47fba1e4a
4 changed files with 309 additions and 80 deletions

View File

@@ -306,3 +306,169 @@ reveal_type(b) # revealed: Unknown
reveal_type(a) # revealed: LiteralString
reveal_type(b) # revealed: LiteralString
```
## Union
### Same types
Union of two tuples of equal length and each element is of the same type.
```py
def _(arg: tuple[int, int] | tuple[int, int]):
(a, b) = arg
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int
```
### Mixed types (1)
Union of two tuples of equal length and one element differs in its type.
```py
def _(arg: tuple[int, int] | tuple[int, str]):
a, b = arg
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int | str
```
### Mixed types (2)
Union of two tuples of equal length and both the element types are different.
```py
def _(arg: tuple[int, str] | tuple[str, int]):
a, b = arg
reveal_type(a) # revealed: int | str
reveal_type(b) # revealed: str | int
```
### Mixed types (3)
Union of three tuples of equal length and various combination of element types:
1. All same types
1. One different type
1. All different types
```py
def _(arg: tuple[int, int, int] | tuple[int, str, bytes] | tuple[int, int, str]):
a, b, c = arg
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int | str
reveal_type(c) # revealed: int | bytes | str
```
### Nested
```py
def _(arg: tuple[int, tuple[str, bytes]] | tuple[tuple[int, bytes], Literal["ab"]]):
a, (b, c) = arg
reveal_type(a) # revealed: int | tuple[int, bytes]
reveal_type(b) # revealed: str
reveal_type(c) # revealed: bytes | LiteralString
```
### Starred expression
```py
def _(arg: tuple[int, bytes, int] | tuple[int, int, str, int, bytes]):
a, *b, c = arg
reveal_type(a) # revealed: int
# TODO: Should be `list[bytes | int | str]`
reveal_type(b) # revealed: @Todo(starred unpacking)
reveal_type(c) # revealed: int | bytes
```
### Size mismatch (1)
```py
def _(arg: tuple[int, bytes, int] | tuple[int, int, str, int, bytes]):
# TODO: Add diagnostic (too many values to unpack)
a, b = arg
reveal_type(a) # revealed: int
reveal_type(b) # revealed: bytes | int
```
### Size mismatch (2)
```py
def _(arg: tuple[int, bytes] | tuple[int, str]):
# TODO: Add diagnostic (there aren't enough values to unpack)
a, b, c = arg
reveal_type(a) # revealed: int
reveal_type(b) # revealed: bytes | str
reveal_type(c) # revealed: Unknown
```
### Same literal types
```py
def _(flag: bool):
if flag:
value = (1, 2)
else:
value = (3, 4)
a, b = value
reveal_type(a) # revealed: Literal[1, 3]
reveal_type(b) # revealed: Literal[2, 4]
```
### Mixed literal types
```py
def _(flag: bool):
if flag:
value = (1, 2)
else:
value = ("a", "b")
a, b = value
reveal_type(a) # revealed: Literal[1] | Literal["a"]
reveal_type(b) # revealed: Literal[2] | Literal["b"]
```
### Typing literal
```py
from typing import Literal
def _(arg: tuple[int, int] | Literal["ab"]):
a, b = arg
reveal_type(a) # revealed: int | LiteralString
reveal_type(b) # revealed: int | LiteralString
```
### Custom iterator (1)
```py
class Iterator:
def __next__(self) -> tuple[int, int] | tuple[int, str]:
return (1, 2)
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
((a, b), c) = Iterable()
reveal_type(a) # revealed: int
reveal_type(b) # revealed: int | str
reveal_type(c) # revealed: tuple[int, int] | tuple[int, str]
```
### Custom iterator (2)
```py
class Iterator:
def __next__(self) -> bytes:
return b""
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
def _(arg: tuple[int, str] | Iterable):
a, b = arg
reveal_type(a) # revealed: int | bytes
reveal_type(b) # revealed: str | bytes
```