mirror of
https://github.com/astral-sh/ruff.git
synced 2025-08-22 11:24:35 +00:00
[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:
parent
ee69d38000
commit
da8aa6a631
3 changed files with 26 additions and 3 deletions
|
@ -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() {
|
||||
|
|
|
@ -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),
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue