[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:
@@ -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
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user