[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:
committed by
GitHub
parent
2a478ce1b2
commit
da6b68cb58
@@ -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
|
||||
|
||||
@@ -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]
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user