mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-23 00:31:55 +00:00
## Summary Previously we held off from doing this because we weren't sure that it was worth the added complexity cost. But our code has changed in the months since we made that initial decision, and I think the structure of the code is such that it no longer really leads to much added complexity to add precise inference when unpacking a string literal or a bytes literal. The improved inference we gain from this has real benefits to users (see the mypy_primer report), and this PR doesn't appear to have a performance impact. ## Test plan mdtests
17 KiB
17 KiB
For loops
Basic for loop
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
for x in IntIterable():
pass
# revealed: int
# error: [possibly-unresolved-reference]
reveal_type(x)
With previous definition
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
x = "foo"
for x in IntIterable():
pass
reveal_type(x) # revealed: Literal["foo"] | int
With else (no break)
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
for x in IntIterable():
pass
else:
x = "foo"
reveal_type(x) # revealed: Literal["foo"]
May break
class IntIterator:
def __next__(self) -> int:
return 42
class IntIterable:
def __iter__(self) -> IntIterator:
return IntIterator()
for x in IntIterable():
if x > 5:
break
else:
x = "foo"
reveal_type(x) # revealed: int | Literal["foo"]
With old-style iteration protocol
class OldStyleIterable:
def __getitem__(self, key: int) -> int:
return 42
for x in OldStyleIterable():
pass
# revealed: int
# error: [possibly-unresolved-reference]
reveal_type(x)
With heterogeneous tuple
for x in (1, "a", b"foo"):
pass
# revealed: Literal[1, "a", b"foo"]
# error: [possibly-unresolved-reference]
reveal_type(x)
With non-callable iterator
from typing_extensions import reveal_type
def _(flag: bool):
class NotIterable:
if flag:
__iter__: int = 1
else:
__iter__: None = None
# error: [not-iterable]
for x in NotIterable():
pass
# revealed: Unknown
# error: [possibly-unresolved-reference]
reveal_type(x)
Invalid iterable
nonsense = 123
for x in nonsense: # error: [not-iterable]
pass
New over old style iteration protocol
class NotIterable:
def __getitem__(self, key: int) -> int:
return 42
__iter__: None = None
for x in NotIterable(): # error: [not-iterable]
pass
Union type as iterable
class TestIter:
def __next__(self) -> int:
return 42
class Test:
def __iter__(self) -> TestIter:
return TestIter()
class Test2:
def __iter__(self) -> TestIter:
return TestIter()
def _(flag: bool):
for x in Test() if flag else Test2():
reveal_type(x) # revealed: int
Union type as iterator
class TestIter:
def __next__(self) -> int:
return 42
class TestIter2:
def __next__(self) -> int:
return 42
class Test:
def __iter__(self) -> TestIter | TestIter2:
return TestIter()
for x in Test():
reveal_type(x) # revealed: int
Union type as iterable and union type as iterator
class Result1A: ...
class Result1B: ...
class Result2A: ...
class Result2B: ...
class Result3: ...
class Result4: ...
class TestIter1:
def __next__(self) -> Result1A | Result1B:
return Result1B()
class TestIter2:
def __next__(self) -> Result2A | Result2B:
return Result2B()
class TestIter3:
def __next__(self) -> Result3:
return Result3()
class TestIter4:
def __next__(self) -> Result4:
return Result4()
class Test:
def __iter__(self) -> TestIter1 | TestIter2:
return TestIter1()
class Test2:
def __iter__(self) -> TestIter3 | TestIter4:
return TestIter3()
def _(flag: bool):
for x in Test() if flag else Test2():
reveal_type(x) # revealed: Result1A | Result1B | Result2A | Result2B | Result3 | Result4
Union type as iterable where one union element has no __iter__ method
from typing_extensions import reveal_type
class TestIter:
def __next__(self) -> int:
return 42
class Test:
def __iter__(self) -> TestIter:
return TestIter()
def _(flag: bool):
# error: [not-iterable]
for x in Test() if flag else 42:
reveal_type(x) # revealed: int
Union type as iterable where one union element has invalid __iter__ method
from typing_extensions import reveal_type
class TestIter:
def __next__(self) -> int:
return 42
class Test:
def __iter__(self) -> TestIter:
return TestIter()
class Test2:
def __iter__(self) -> int:
return 42
def _(flag: bool):
# TODO: Improve error message to state which union variant isn't iterable (https://github.com/astral-sh/ruff/issues/13989)
# error: [not-iterable]
for x in Test() if flag else Test2():
reveal_type(x) # revealed: int
Union type as iterator where one union element has no __next__ method
class TestIter:
def __next__(self) -> int:
return 42
class Test:
def __iter__(self) -> TestIter | int:
return TestIter()
# error: [not-iterable] "Object of type `Test` may not be iterable"
for x in Test():
reveal_type(x) # revealed: int
Possibly-not-callable __iter__ method
def _(flag: bool):
class Iterator:
def __next__(self) -> int:
return 42
class CustomCallable:
if flag:
def __call__(self, *args, **kwargs) -> Iterator:
return Iterator()
else:
__call__: None = None
class Iterable1:
__iter__: CustomCallable = CustomCallable()
class Iterable2:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
else:
__iter__: None = None
# error: [not-iterable] "Object of type `Iterable1` may not be iterable"
for x in Iterable1():
# TODO... `int` might be ideal here?
reveal_type(x) # revealed: int | Unknown
# error: [not-iterable] "Object of type `Iterable2` may not be iterable"
for y in Iterable2():
# TODO... `int` might be ideal here?
reveal_type(y) # revealed: int | Unknown
__iter__ method with a bad signature
from typing_extensions import reveal_type
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self, extra_arg) -> Iterator:
return Iterator()
# error: [not-iterable]
for x in Iterable():
reveal_type(x) # revealed: int
__iter__ does not return an iterator
from typing_extensions import reveal_type
class Bad:
def __iter__(self) -> int:
return 42
# error: [not-iterable]
for x in Bad():
reveal_type(x) # revealed: Unknown
__iter__ returns an object with a possibly unbound __next__ method
def _(flag: bool):
class Iterator:
if flag:
def __next__(self) -> int:
return 42
class Iterable:
def __iter__(self) -> Iterator:
return Iterator()
# error: [not-iterable] "Object of type `Iterable` may not be iterable"
for x in Iterable():
reveal_type(x) # revealed: int
__iter__ returns an iterator with an invalid __next__ method
from typing_extensions import reveal_type
class Iterator1:
def __next__(self, extra_arg) -> int:
return 42
class Iterator2:
__next__: None = None
class Iterable1:
def __iter__(self) -> Iterator1:
return Iterator1()
class Iterable2:
def __iter__(self) -> Iterator2:
return Iterator2()
# error: [not-iterable]
for x in Iterable1():
reveal_type(x) # revealed: int
# error: [not-iterable]
for y in Iterable2():
reveal_type(y) # revealed: Unknown
Possibly unbound __iter__ and bad __getitem__ method
from typing_extensions import reveal_type
def _(flag: bool):
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
# invalid signature because it only accepts a `str`,
# but the old-style iteration protocol will pass it an `int`
def __getitem__(self, key: str) -> bytes:
return bytes()
# error: [not-iterable]
for x in Iterable():
reveal_type(x) # revealed: int | bytes
Possibly unbound __iter__ and not-callable __getitem__
This snippet tests that we infer the element type correctly in the following edge case:
__iter__is a method with the correct parameter spec that returns a valid iterator; BUT__iter__is possibly unbound; AND__getitem__is set to a non-callable type
It's important that we emit a diagnostic here, but it's also important that we still use the return
type of the iterator's __next__ method as the inferred type of x in the for loop:
def _(flag: bool):
class Iterator:
def __next__(self) -> int:
return 42
class Iterable:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
__getitem__: None = None
# error: [not-iterable] "Object of type `Iterable` may not be iterable"
for x in Iterable():
reveal_type(x) # revealed: int
Possibly unbound __iter__ and possibly unbound __getitem__
from typing_extensions import reveal_type
class Iterator:
def __next__(self) -> int:
return 42
def _(flag1: bool, flag2: bool):
class Iterable:
if flag1:
def __iter__(self) -> Iterator:
return Iterator()
if flag2:
def __getitem__(self, key: int) -> bytes:
return bytes()
# error: [not-iterable]
for x in Iterable():
reveal_type(x) # revealed: int | bytes
No __iter__ method and __getitem__ is not callable
from typing_extensions import reveal_type
class Bad:
__getitem__: None = None
# error: [not-iterable]
for x in Bad():
reveal_type(x) # revealed: Unknown
Possibly-not-callable __getitem__ method
from typing_extensions import reveal_type
def _(flag: bool):
class CustomCallable:
if flag:
def __call__(self, *args, **kwargs) -> int:
return 42
else:
__call__: None = None
class Iterable1:
__getitem__: CustomCallable = CustomCallable()
class Iterable2:
if flag:
def __getitem__(self, key: int) -> int:
return 42
else:
__getitem__: None = None
# error: [not-iterable]
for x in Iterable1():
# TODO... `int` might be ideal here?
reveal_type(x) # revealed: int | Unknown
# error: [not-iterable]
for y in Iterable2():
# TODO... `int` might be ideal here?
reveal_type(y) # revealed: int | Unknown
Bad __getitem__ method
from typing_extensions import reveal_type
class Iterable:
# invalid because it will implicitly be passed an `int`
# by the interpreter
def __getitem__(self, key: str) -> int:
return 42
# error: [not-iterable]
for x in Iterable():
reveal_type(x) # revealed: int
Possibly unbound __iter__ but definitely bound __getitem__
Here, we should not emit a diagnostic: if __iter__ is unbound, we should fallback to
__getitem__:
class Iterator:
def __next__(self) -> str:
return "foo"
def _(flag: bool):
class Iterable:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
def __getitem__(self, key: int) -> bytes:
return b"foo"
for x in Iterable():
reveal_type(x) # revealed: str | bytes
Possibly invalid __iter__ methods
from typing_extensions import reveal_type
class Iterator:
def __next__(self) -> int:
return 42
def _(flag: bool):
class Iterable1:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
else:
def __iter__(self, invalid_extra_arg) -> Iterator:
return Iterator()
# error: [not-iterable]
for x in Iterable1():
reveal_type(x) # revealed: int
class Iterable2:
if flag:
def __iter__(self) -> Iterator:
return Iterator()
else:
__iter__: None = None
# error: [not-iterable]
for x in Iterable2():
# TODO: `int` would probably be better here:
reveal_type(x) # revealed: int | Unknown
Possibly invalid __next__ method
from typing_extensions import reveal_type
def _(flag: bool):
class Iterator1:
if flag:
def __next__(self) -> int:
return 42
else:
def __next__(self, invalid_extra_arg) -> str:
return "foo"
class Iterator2:
if flag:
def __next__(self) -> int:
return 42
else:
__next__: None = None
class Iterable1:
def __iter__(self) -> Iterator1:
return Iterator1()
class Iterable2:
def __iter__(self) -> Iterator2:
return Iterator2()
# error: [not-iterable]
for x in Iterable1():
reveal_type(x) # revealed: int | str
# error: [not-iterable]
for y in Iterable2():
# TODO: `int` would probably be better here:
reveal_type(y) # revealed: int | Unknown
Possibly invalid __getitem__ methods
from typing_extensions import reveal_type
def _(flag: bool):
class Iterable1:
if flag:
def __getitem__(self, item: int) -> str:
return "foo"
else:
__getitem__: None = None
class Iterable2:
if flag:
def __getitem__(self, item: int) -> str:
return "foo"
else:
def __getitem__(self, item: str) -> int:
return 42
# error: [not-iterable]
for x in Iterable1():
# TODO: `str` might be better
reveal_type(x) # revealed: str | Unknown
# error: [not-iterable]
for y in Iterable2():
reveal_type(y) # revealed: str | int
Possibly unbound __iter__ and possibly invalid __getitem__
from typing_extensions import reveal_type
class Iterator:
def __next__(self) -> bytes:
return b"foo"
def _(flag: bool, flag2: bool):
class Iterable1:
if flag:
def __getitem__(self, item: int) -> str:
return "foo"
else:
__getitem__: None = None
if flag2:
def __iter__(self) -> Iterator:
return Iterator()
class Iterable2:
if flag:
def __getitem__(self, item: int) -> str:
return "foo"
else:
def __getitem__(self, item: str) -> int:
return 42
if flag2:
def __iter__(self) -> Iterator:
return Iterator()
# error: [not-iterable]
for x in Iterable1():
# TODO: `bytes | str` might be better
reveal_type(x) # revealed: bytes | str | Unknown
# error: [not-iterable]
for y in Iterable2():
reveal_type(y) # revealed: bytes | str | int
Empty tuple is iterable
for x in ():
reveal_type(x) # revealed: Never
Never is iterable
from typing_extensions import Never
def f(never: Never):
for x in never:
reveal_type(x) # revealed: Unknown
Iterating over literals
from typing import Literal
for char in "abcde":
reveal_type(char) # revealed: Literal["a", "b", "c", "d", "e"]
for char in b"abcde":
reveal_type(char) # revealed: Literal[97, 98, 99, 100, 101]
A class literal is iterable if it inherits from Any
A class literal can be iterated over if it has Any or Unknown in its MRO, since the
Any/Unknown element in the MRO could materialize to a class with a custom metaclass that defines
__iter__ for all instances of the metaclass:
from unresolved_module import SomethingUnknown # error: [unresolved-import]
from typing import Any, Iterable
from ty_extensions import static_assert, is_assignable_to, TypeOf, Unknown
class Foo(SomethingUnknown): ...
reveal_type(Foo.__mro__) # revealed: tuple[<class 'Foo'>, Unknown, <class 'object'>]
# TODO: these should pass
static_assert(is_assignable_to(TypeOf[Foo], Iterable[Unknown])) # error: [static-assert-error]
static_assert(is_assignable_to(type[Foo], Iterable[Unknown])) # error: [static-assert-error]
# TODO: should not error
# error: [not-iterable]
for x in Foo:
reveal_type(x) # revealed: Unknown
class Bar(Any): ...
reveal_type(Bar.__mro__) # revealed: tuple[<class 'Bar'>, Any, <class 'object'>]
# TODO: these should pass
static_assert(is_assignable_to(TypeOf[Bar], Iterable[Any])) # error: [static-assert-error]
static_assert(is_assignable_to(type[Bar], Iterable[Any])) # error: [static-assert-error]
# TODO: should not error
# error: [not-iterable]
for x in Bar:
# TODO: should reveal `Any`
reveal_type(x) # revealed: Unknown