[ty] Use range instead of custom IntIterable (#21138)

## Summary

We previously didn't understand `range` and wrote these custom
`IntIterable`/`IntIterator` classes for tests. We can now remove them
and make the tests shorter in some places.
This commit is contained in:
David Peter 2025-10-30 15:21:55 +01:00 committed by GitHub
parent 1ebedf6df5
commit 1b0ee4677e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 35 additions and 115 deletions

View file

@ -300,14 +300,6 @@ reveal_type(c_instance.b) # revealed: Unknown | list[Literal[2, 3]]
#### Attributes defined in for-loop (unpacking)
```py
class IntIterator:
def __next__(self) -> int:
return 1
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
class TupleIterator:
def __next__(self) -> tuple[int, str]:
return (1, "a")
@ -320,7 +312,7 @@ class NonIterable: ...
class C:
def __init__(self):
for self.x in IntIterable():
for self.x in range(3):
pass
for _, self.y in TupleIterable():
@ -378,14 +370,6 @@ reveal_type(c_instance.y) # revealed: Unknown | int
#### Attributes defined in comprehensions
```py
class IntIterator:
def __next__(self) -> int:
return 1
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
class TupleIterator:
def __next__(self) -> tuple[int, str]:
return (1, "a")
@ -398,7 +382,7 @@ class C:
def __init__(self) -> None:
# TODO: Should not emit this diagnostic
# error: [unresolved-attribute]
[... for self.a in IntIterable()]
[... for self.a in range(3)]
# TODO: Should not emit this diagnostic
# error: [unresolved-attribute]
# error: [unresolved-attribute]
@ -406,11 +390,11 @@ class C:
# TODO: Should not emit this diagnostic
# error: [unresolved-attribute]
# error: [unresolved-attribute]
[... for self.d in IntIterable() for self.e in IntIterable()]
[... for self.d in range(3) for self.e in range(3)]
# TODO: Should not emit this diagnostic
# error: [unresolved-attribute]
[[... for self.f in IntIterable()] for _ in IntIterable()]
[[... for self.g in IntIterable()] for self in [D()]]
[[... for self.f in range(3)] for _ in range(3)]
[[... for self.g in range(3)] for self in [D()]]
class D:
g: int
@ -2058,16 +2042,8 @@ mod.global_symbol = 1
# TODO: this should be an error, but we do not understand list unpackings yet.
[_, mod.global_symbol] = [1, 2]
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
# error: [invalid-assignment] "Object of type `int` is not assignable to attribute `global_symbol` of type `str`"
for mod.global_symbol in IntIterable():
for mod.global_symbol in range(3):
pass
```

View file

@ -3,71 +3,47 @@
## Basic comprehensions
```py
class IntIterator:
def __next__(self) -> int:
return 42
# revealed: int
[reveal_type(x) for x in range(3)]
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
class Row:
def __next__(self) -> range:
return range(3)
class Table:
def __iter__(self) -> Row:
return Row()
# revealed: tuple[int, range]
[reveal_type((cell, row)) for row in Table() for cell in row]
# revealed: int
[reveal_type(x) for x in IntIterable()]
class IteratorOfIterables:
def __next__(self) -> IntIterable:
return IntIterable()
class IterableOfIterables:
def __iter__(self) -> IteratorOfIterables:
return IteratorOfIterables()
# revealed: tuple[int, IntIterable]
[reveal_type((x, y)) for y in IterableOfIterables() for x in y]
{reveal_type(x): 0 for x in range(3)}
# revealed: int
{reveal_type(x): 0 for x in IntIterable()}
# revealed: int
{0: reveal_type(x) for x in IntIterable()}
{0: reveal_type(x) for x in range(3)}
```
## Nested comprehension
```py
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
# revealed: tuple[int, int]
[[reveal_type((x, y)) for x in IntIterable()] for y in IntIterable()]
[[reveal_type((x, y)) for x in range(3)] for y in range(3)]
```
## Comprehension referencing outer comprehension
```py
class IntIterator:
def __next__(self) -> int:
return 42
class Row:
def __next__(self) -> range:
return range(3)
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
class Table:
def __iter__(self) -> Row:
return Row()
class IteratorOfIterables:
def __next__(self) -> IntIterable:
return IntIterable()
class IterableOfIterables:
def __iter__(self) -> IteratorOfIterables:
return IteratorOfIterables()
# revealed: tuple[int, IntIterable]
[[reveal_type((x, y)) for x in y] for y in IterableOfIterables()]
# revealed: tuple[int, range]
[[reveal_type((cell, row)) for cell in row] for row in Table()]
```
## Comprehension with unbound iterable
@ -79,17 +55,9 @@ Iterating over an unbound iterable yields `Unknown`:
# revealed: Unknown
[reveal_type(z) for z in x]
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
# error: [not-iterable] "Object of type `int` is not iterable"
# revealed: tuple[int, Unknown]
[reveal_type((x, z)) for x in IntIterable() for z in x]
[reveal_type((x, z)) for x in range(3) for z in x]
```
## Starred expressions
@ -99,16 +67,8 @@ Starred expressions must be iterable
```py
class NotIterable: ...
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
# This is fine:
x = [*Iterable()]
x = [*range(3)]
# error: [not-iterable] "Object of type `NotIterable` is not iterable"
y = [*NotIterable()]
@ -138,16 +98,8 @@ This tests that we understand that `async` comprehensions do *not* work accordin
iteration protocol
```py
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
async def _():
# error: [not-iterable] "Object of type `Iterable` is not async-iterable"
# error: [not-iterable] "Object of type `range` is not async-iterable"
# revealed: Unknown
[reveal_type(x) async for x in Iterable()]
[reveal_type(x) async for x in range(3)]
```

View file

@ -1,14 +1,6 @@
# Comprehensions with invalid syntax
```py
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
# Missing 'in' keyword.
# It's reasonably clear here what they *meant* to write,
@ -16,7 +8,7 @@ class IntIterable:
# error: [invalid-syntax] "Expected 'in', found name"
# revealed: int
[reveal_type(a) for a IntIterable()]
[reveal_type(a) for a range(3)]
# Missing iteration variable
@ -25,7 +17,7 @@ class IntIterable:
# error: [invalid-syntax] "Expected 'in', found name"
# error: [unresolved-reference]
# revealed: Unknown
[reveal_type(b) for in IntIterable()]
[reveal_type(b) for in range(3)]
# Missing iterable