From da8aa6a6312ee198489f974e2d7b665a7699b7ab Mon Sep 17 00:00:00 2001 From: David Peter Date: Tue, 22 Jul 2025 16:09:28 +0200 Subject: [PATCH] [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 https://github.com/astral-sh/ty/issues/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 https://github.com/astral-sh/ty/issues/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. --- .../ty_python_semantic/resources/mdtest/enums.md | 3 +-- crates/ty_python_semantic/src/types/call/bind.rs | 14 ++++++++++++++ crates/ty_python_semantic/src/types/class.rs | 12 +++++++++++- 3 files changed, 26 insertions(+), 3 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 482fdd4960..e804e984a8 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -406,8 +406,7 @@ class Color(Enum): BLUE = 3 for color in Color: - # TODO: Should be `Color` - reveal_type(color) # revealed: Unknown + reveal_type(color) # revealed: Color # TODO: Should be `list[Color]` reveal_type(list(Color)) # revealed: list[Unknown] diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 1cb121e3ed..3bbda82594 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -19,6 +19,7 @@ use crate::types::diagnostic::{ NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS, UNKNOWN_ARGUMENT, }; +use crate::types::enums::is_enum_class; use crate::types::function::{ DataclassTransformerParams, FunctionDecorators, FunctionType, KnownFunction, OverloadLiteral, }; @@ -560,6 +561,19 @@ impl<'db> Bindings<'db> { } } + // TODO: This branch can be removed once https://github.com/astral-sh/ty/issues/501 is resolved + Type::BoundMethod(bound_method) + if bound_method.function(db).name(db) == "__iter__" + && is_enum_class(db, bound_method.self_instance(db)) => + { + if let Some(enum_instance) = bound_method.self_instance(db).to_instance(db) + { + overload.set_return_type( + KnownClass::Iterator.to_specialized_instance(db, [enum_instance]), + ); + } + } + Type::FunctionLiteral(function_type) => match function_type.known(db) { Some(KnownFunction::IsEquivalentTo) => { if let [Some(ty_a), Some(ty_b)] = overload.parameter_types() { diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index f4f21cfab2..34ba5797bd 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -2564,6 +2564,7 @@ pub enum KnownClass { NewType, SupportsIndex, Iterable, + Iterator, // Collections ChainMap, Counter, @@ -2660,6 +2661,7 @@ impl KnownClass { | Self::Nonmember | Self::ABCMeta | Self::Iterable + | Self::Iterator // Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue | Self::NamedTuple // Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9 @@ -2753,6 +2755,7 @@ impl KnownClass { | Self::OrderedDict | Self::NewType | Self::Iterable + | Self::Iterator | Self::BaseExceptionGroup => false, } } @@ -2814,6 +2817,7 @@ impl KnownClass { | KnownClass::NewType | KnownClass::SupportsIndex | KnownClass::Iterable + | KnownClass::Iterator | KnownClass::ChainMap | KnownClass::Counter | KnownClass::DefaultDict @@ -2842,7 +2846,7 @@ impl KnownClass { /// 2. It's probably more performant. const fn is_protocol(self) -> bool { match self { - Self::SupportsIndex | Self::Iterable => true, + Self::SupportsIndex | Self::Iterable | Self::Iterator => true, Self::Any | Self::Bool @@ -2968,6 +2972,7 @@ impl KnownClass { Self::ABCMeta => "ABCMeta", Self::Super => "super", Self::Iterable => "Iterable", + Self::Iterator => "Iterator", // For example, `typing.List` is defined as `List = _Alias()` in typeshed Self::StdlibAlias => "_Alias", // This is the name the type of `sys.version_info` has in typeshed, @@ -3203,6 +3208,7 @@ impl KnownClass { | Self::NamedTuple | Self::StdlibAlias | Self::Iterable + | Self::Iterator | Self::SupportsIndex => KnownModule::Typing, Self::TypeAliasType | Self::TypeVarTuple @@ -3311,6 +3317,7 @@ impl KnownClass { | Self::Field | Self::KwOnly | Self::Iterable + | Self::Iterator | Self::NamedTupleFallback => false, } } @@ -3384,6 +3391,7 @@ impl KnownClass { | Self::Field | Self::KwOnly | Self::Iterable + | Self::Iterator | Self::NamedTupleFallback => false, } } @@ -3435,6 +3443,7 @@ impl KnownClass { "TypeAliasType" => Self::TypeAliasType, "TypeVar" => Self::TypeVar, "Iterable" => Self::Iterable, + "Iterator" => Self::Iterator, "ParamSpec" => Self::ParamSpec, "ParamSpecArgs" => Self::ParamSpecArgs, "ParamSpecKwargs" => Self::ParamSpecKwargs, @@ -3538,6 +3547,7 @@ impl KnownClass { | Self::TypeVarTuple | Self::NamedTuple | Self::Iterable + | Self::Iterator | Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions), Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),