Skip to content

Commit da8aa6a

Browse files
authored
[ty] Support iterating over enums (#19486)
## Summary Infer the correct type in a scenario like this: ```py class Color(Enum): RED = 1 GREEN = 2 BLUE = 3 for color in Color: reveal_type(color) # revealed: Color ``` We should eventually support this out-of-the-box when astral-sh/ty#501 is implemented. For this reason, @AlexWaygood would prefer to keep things as they are (we currently infer `Unknown`, so false positives seem unlikely). But it seemed relatively easy to support, so I'm opening this for discussion. part of astral-sh/ty#183 ## Test Plan Adapted existing test. ## Ecosystem analysis ```diff - warning[unused-ignore-comment] rotkehlchen/chain/aggregator.py:591:82: Unused blanket `type: ignore` directive ``` This `unused-ignore-comment` goes away due to a new true positive.
1 parent ee69d38 commit da8aa6a

3 files changed

Lines changed: 26 additions & 3 deletions

File tree

crates/ty_python_semantic/resources/mdtest/enums.md

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -406,8 +406,7 @@ class Color(Enum):
406406
BLUE = 3
407407

408408
for color in Color:
409-
# TODO: Should be `Color`
410-
reveal_type(color) # revealed: Unknown
409+
reveal_type(color) # revealed: Color
411410

412411
# TODO: Should be `list[Color]`
413412
reveal_type(list(Color)) # revealed: list[Unknown]

crates/ty_python_semantic/src/types/call/bind.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ use crate::types::diagnostic::{
1919
NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
2020
UNKNOWN_ARGUMENT,
2121
};
22+
use crate::types::enums::is_enum_class;
2223
use crate::types::function::{
2324
DataclassTransformerParams, FunctionDecorators, FunctionType, KnownFunction, OverloadLiteral,
2425
};
@@ -560,6 +561,19 @@ impl<'db> Bindings<'db> {
560561
}
561562
}
562563

564+
// TODO: This branch can be removed once https://github.com/astral-sh/ty/issues/501 is resolved
565+
Type::BoundMethod(bound_method)
566+
if bound_method.function(db).name(db) == "__iter__"
567+
&& is_enum_class(db, bound_method.self_instance(db)) =>
568+
{
569+
if let Some(enum_instance) = bound_method.self_instance(db).to_instance(db)
570+
{
571+
overload.set_return_type(
572+
KnownClass::Iterator.to_specialized_instance(db, [enum_instance]),
573+
);
574+
}
575+
}
576+
563577
Type::FunctionLiteral(function_type) => match function_type.known(db) {
564578
Some(KnownFunction::IsEquivalentTo) => {
565579
if let [Some(ty_a), Some(ty_b)] = overload.parameter_types() {

crates/ty_python_semantic/src/types/class.rs

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2564,6 +2564,7 @@ pub enum KnownClass {
25642564
NewType,
25652565
SupportsIndex,
25662566
Iterable,
2567+
Iterator,
25672568
// Collections
25682569
ChainMap,
25692570
Counter,
@@ -2660,6 +2661,7 @@ impl KnownClass {
26602661
| Self::Nonmember
26612662
| Self::ABCMeta
26622663
| Self::Iterable
2664+
| Self::Iterator
26632665
// Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
26642666
| Self::NamedTuple
26652667
// Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
@@ -2753,6 +2755,7 @@ impl KnownClass {
27532755
| Self::OrderedDict
27542756
| Self::NewType
27552757
| Self::Iterable
2758+
| Self::Iterator
27562759
| Self::BaseExceptionGroup => false,
27572760
}
27582761
}
@@ -2814,6 +2817,7 @@ impl KnownClass {
28142817
| KnownClass::NewType
28152818
| KnownClass::SupportsIndex
28162819
| KnownClass::Iterable
2820+
| KnownClass::Iterator
28172821
| KnownClass::ChainMap
28182822
| KnownClass::Counter
28192823
| KnownClass::DefaultDict
@@ -2842,7 +2846,7 @@ impl KnownClass {
28422846
/// 2. It's probably more performant.
28432847
const fn is_protocol(self) -> bool {
28442848
match self {
2845-
Self::SupportsIndex | Self::Iterable => true,
2849+
Self::SupportsIndex | Self::Iterable | Self::Iterator => true,
28462850

28472851
Self::Any
28482852
| Self::Bool
@@ -2968,6 +2972,7 @@ impl KnownClass {
29682972
Self::ABCMeta => "ABCMeta",
29692973
Self::Super => "super",
29702974
Self::Iterable => "Iterable",
2975+
Self::Iterator => "Iterator",
29712976
// For example, `typing.List` is defined as `List = _Alias()` in typeshed
29722977
Self::StdlibAlias => "_Alias",
29732978
// This is the name the type of `sys.version_info` has in typeshed,
@@ -3203,6 +3208,7 @@ impl KnownClass {
32033208
| Self::NamedTuple
32043209
| Self::StdlibAlias
32053210
| Self::Iterable
3211+
| Self::Iterator
32063212
| Self::SupportsIndex => KnownModule::Typing,
32073213
Self::TypeAliasType
32083214
| Self::TypeVarTuple
@@ -3311,6 +3317,7 @@ impl KnownClass {
33113317
| Self::Field
33123318
| Self::KwOnly
33133319
| Self::Iterable
3320+
| Self::Iterator
33143321
| Self::NamedTupleFallback => false,
33153322
}
33163323
}
@@ -3384,6 +3391,7 @@ impl KnownClass {
33843391
| Self::Field
33853392
| Self::KwOnly
33863393
| Self::Iterable
3394+
| Self::Iterator
33873395
| Self::NamedTupleFallback => false,
33883396
}
33893397
}
@@ -3435,6 +3443,7 @@ impl KnownClass {
34353443
"TypeAliasType" => Self::TypeAliasType,
34363444
"TypeVar" => Self::TypeVar,
34373445
"Iterable" => Self::Iterable,
3446+
"Iterator" => Self::Iterator,
34383447
"ParamSpec" => Self::ParamSpec,
34393448
"ParamSpecArgs" => Self::ParamSpecArgs,
34403449
"ParamSpecKwargs" => Self::ParamSpecKwargs,
@@ -3538,6 +3547,7 @@ impl KnownClass {
35383547
| Self::TypeVarTuple
35393548
| Self::NamedTuple
35403549
| Self::Iterable
3550+
| Self::Iterator
35413551
| Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
35423552
Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),
35433553

0 commit comments

Comments
 (0)