[red-knot] infer attribute assignments bound in comprehensions (#17396)

## Summary

This PR is a follow-up to #16852.

Instance variables bound in comprehensions are recorded, allowing type
inference to work correctly.

This required adding support for unpacking in comprehension which
resolves https://github.com/astral-sh/ruff/issues/15369.

## Test Plan

One TODO in `mdtest/attributes.md` is now resolved, and some new test
cases are added.

---------

Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>
This commit is contained in:
Shunsuke Shibayama
2025-04-19 10:12:48 +09:00
committed by GitHub
parent 2a478ce1b2
commit da6b68cb58
10 changed files with 349 additions and 108 deletions

View File

@@ -397,15 +397,27 @@ class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
class TupleIterator:
def __next__(self) -> tuple[int, str]:
return (1, "a")
class TupleIterable:
def __iter__(self) -> TupleIterator:
return TupleIterator()
class C:
def __init__(self) -> None:
[... for self.a in IntIterable()]
[... for (self.b, self.c) in TupleIterable()]
[... for self.d in IntIterable() for self.e in IntIterable()]
c_instance = C()
# TODO: Should be `Unknown | int`
# error: [unresolved-attribute]
reveal_type(c_instance.a) # revealed: Unknown
reveal_type(c_instance.a) # revealed: Unknown | int
reveal_type(c_instance.b) # revealed: Unknown | int
reveal_type(c_instance.c) # revealed: Unknown | str
reveal_type(c_instance.d) # revealed: Unknown | int
reveal_type(c_instance.e) # revealed: Unknown | int
```
#### Conditionally declared / bound attributes

View File

@@ -708,3 +708,95 @@ with ContextManager() as (a, b, c):
reveal_type(b) # revealed: Unknown
reveal_type(c) # revealed: Unknown
```
## Comprehension
Unpacking in a comprehension.
### Same types
```py
def _(arg: tuple[tuple[int, int], tuple[int, int]]):
# revealed: tuple[int, int]
[reveal_type((a, b)) for a, b in arg]
```
### Mixed types (1)
```py
def _(arg: tuple[tuple[int, int], tuple[int, str]]):
# revealed: tuple[int, int | str]
[reveal_type((a, b)) for a, b in arg]
```
### Mixed types (2)
```py
def _(arg: tuple[tuple[int, str], tuple[str, int]]):
# revealed: tuple[int | str, str | int]
[reveal_type((a, b)) for a, b in arg]
```
### Mixed types (3)
```py
def _(arg: tuple[tuple[int, int, int], tuple[int, str, bytes], tuple[int, int, str]]):
# revealed: tuple[int, int | str, int | bytes | str]
[reveal_type((a, b, c)) for a, b, c in arg]
```
### Same literal values
```py
# revealed: tuple[Literal[1, 3], Literal[2, 4]]
[reveal_type((a, b)) for a, b in ((1, 2), (3, 4))]
```
### Mixed literal values (1)
```py
# revealed: tuple[Literal[1, "a"], Literal[2, "b"]]
[reveal_type((a, b)) for a, b in ((1, 2), ("a", "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"
# error: [invalid-assignment] "Not enough values to unpack (expected 2, got 1)"
# revealed: tuple[Unknown | Literal[3, 5], Unknown | Literal["a", "b"]]
[reveal_type((a, b)) for a, b in (1, 2, (3, "a"), 4, (5, "b"), "c")]
```
### Custom iterator (1)
```py
class Iterator:
def __next__(self) -> tuple[int, int]:
return (1, 2)
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
# revealed: tuple[int, int]
[reveal_type((a, b)) for a, b in Iterable()]
```
### 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]):
# revealed: tuple[int | bytes, str | bytes]
[reveal_type((a, b)) for a, b in arg]
```