ruff/crates/red_knot_python_semantic/resources/mdtest/unpacking.md
Dhruv Manilawala 113c804a62
[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.
2024-12-23 06:13:49 +00:00

13 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)
[a, *b, c, d] = (1, 2)
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(starred unpacking)
reveal_type(c)  # revealed: Literal[2]
reveal_type(d)  # revealed: Unknown

Starred expression (2)

[a, *b, c] = (1, 2)
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(starred unpacking)
reveal_type(c)  # revealed: Literal[2]

Starred expression (3)

[a, *b, c] = (1, 2, 3)
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(starred unpacking)
reveal_type(c)  # revealed: Literal[3]

Starred expression (4)

[a, *b, c, d] = (1, 2, 3, 4, 5, 6)
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(starred unpacking)
reveal_type(c)  # revealed: Literal[5]
reveal_type(d)  # revealed: Literal[6]

Starred expression (5)

[a, b, *c] = (1, 2, 3, 4)
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(starred unpacking)

Starred expression (6)

# TODO: Add diagnostic (need more values to unpack)
(a, b, c, *d, e, f) = (1,)
reveal_type(a)  # revealed: Literal[1]
reveal_type(b)  # revealed: Unknown
reveal_type(c)  # revealed: Unknown
reveal_type(d)  # revealed: @Todo(starred unpacking)
reveal_type(e)  # revealed: Unknown
reveal_type(f)  # revealed: Unknown

Non-iterable unpacking

# 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)
(a, *b, c, d) = "ab"
reveal_type(a)  # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b)  # revealed: @Todo(starred unpacking)
reveal_type(c)  # revealed: LiteralString
reveal_type(d)  # revealed: Unknown

Starred expression (2)

(a, *b, c) = "ab"
reveal_type(a)  # revealed: LiteralString
# TODO: Should be list[Any] once support for assigning to starred expression is added
reveal_type(b)  # revealed: @Todo(starred unpacking)
reveal_type(c)  # revealed: LiteralString

Starred expression (3)

(a, *b, c) = "abc"
reveal_type(a)  # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b)  # revealed: @Todo(starred unpacking)
reveal_type(c)  # revealed: LiteralString

Starred expression (4)

(a, *b, c, d) = "abcdef"
reveal_type(a)  # revealed: LiteralString
# TODO: Should be list[LiteralString] once support for assigning to starred expression is added
reveal_type(b)  # revealed: @Todo(starred unpacking)
reveal_type(c)  # revealed: LiteralString
reveal_type(d)  # revealed: LiteralString

Starred expression (5)

(a, b, *c) = "abcd"
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(starred unpacking)

Unicode

# TODO: Add diagnostic (need more values to unpack)
(a, b) = "é"

reveal_type(a)  # revealed: LiteralString
reveal_type(b)  # revealed: Unknown

Unicode escape (1)

# TODO: Add diagnostic (need more values to unpack)
(a, b) = "\u9E6C"

reveal_type(a)  # revealed: LiteralString
reveal_type(b)  # revealed: Unknown

Unicode escape (2)

# TODO: Add diagnostic (need more values to unpack)
(a, b) = "\U0010FFFF"

reveal_type(a)  # revealed: LiteralString
reveal_type(b)  # revealed: Unknown

Surrogates

(a, b) = "\uD800\uDFFF"

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.

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.

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.

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
  2. One different type
  3. All different types
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

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

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)

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)

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

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

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

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)

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)

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

For statement

Unpacking in a for statement.

Same types

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)

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)

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)

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

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)

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)

# 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)

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)

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