[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.
This commit is contained in:
David Peter 2025-07-22 16:09:28 +02:00 committed by GitHub
parent ee69d38000
commit da8aa6a631
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 26 additions and 3 deletions

View file

@ -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() {

View file

@ -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),