mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-28 18:53:25 +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
|
|
@ -406,8 +406,7 @@ class Color(Enum):
|
||||||
BLUE = 3
|
BLUE = 3
|
||||||
|
|
||||||
for color in Color:
|
for color in Color:
|
||||||
# TODO: Should be `Color`
|
reveal_type(color) # revealed: Color
|
||||||
reveal_type(color) # revealed: Unknown
|
|
||||||
|
|
||||||
# TODO: Should be `list[Color]`
|
# TODO: Should be `list[Color]`
|
||||||
reveal_type(list(Color)) # revealed: list[Unknown]
|
reveal_type(list(Color)) # revealed: list[Unknown]
|
||||||
|
|
|
||||||
|
|
@ -19,6 +19,7 @@ use crate::types::diagnostic::{
|
||||||
NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
|
NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
|
||||||
UNKNOWN_ARGUMENT,
|
UNKNOWN_ARGUMENT,
|
||||||
};
|
};
|
||||||
|
use crate::types::enums::is_enum_class;
|
||||||
use crate::types::function::{
|
use crate::types::function::{
|
||||||
DataclassTransformerParams, FunctionDecorators, FunctionType, KnownFunction, OverloadLiteral,
|
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) {
|
Type::FunctionLiteral(function_type) => match function_type.known(db) {
|
||||||
Some(KnownFunction::IsEquivalentTo) => {
|
Some(KnownFunction::IsEquivalentTo) => {
|
||||||
if let [Some(ty_a), Some(ty_b)] = overload.parameter_types() {
|
if let [Some(ty_a), Some(ty_b)] = overload.parameter_types() {
|
||||||
|
|
|
||||||
|
|
@ -2564,6 +2564,7 @@ pub enum KnownClass {
|
||||||
NewType,
|
NewType,
|
||||||
SupportsIndex,
|
SupportsIndex,
|
||||||
Iterable,
|
Iterable,
|
||||||
|
Iterator,
|
||||||
// Collections
|
// Collections
|
||||||
ChainMap,
|
ChainMap,
|
||||||
Counter,
|
Counter,
|
||||||
|
|
@ -2660,6 +2661,7 @@ impl KnownClass {
|
||||||
| Self::Nonmember
|
| Self::Nonmember
|
||||||
| Self::ABCMeta
|
| Self::ABCMeta
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
// Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
|
// Empty tuples are AlwaysFalse; non-empty tuples are AlwaysTrue
|
||||||
| Self::NamedTuple
|
| Self::NamedTuple
|
||||||
// Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
|
// Evaluating `NotImplementedType` in a boolean context was deprecated in Python 3.9
|
||||||
|
|
@ -2753,6 +2755,7 @@ impl KnownClass {
|
||||||
| Self::OrderedDict
|
| Self::OrderedDict
|
||||||
| Self::NewType
|
| Self::NewType
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
| Self::BaseExceptionGroup => false,
|
| Self::BaseExceptionGroup => false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -2814,6 +2817,7 @@ impl KnownClass {
|
||||||
| KnownClass::NewType
|
| KnownClass::NewType
|
||||||
| KnownClass::SupportsIndex
|
| KnownClass::SupportsIndex
|
||||||
| KnownClass::Iterable
|
| KnownClass::Iterable
|
||||||
|
| KnownClass::Iterator
|
||||||
| KnownClass::ChainMap
|
| KnownClass::ChainMap
|
||||||
| KnownClass::Counter
|
| KnownClass::Counter
|
||||||
| KnownClass::DefaultDict
|
| KnownClass::DefaultDict
|
||||||
|
|
@ -2842,7 +2846,7 @@ impl KnownClass {
|
||||||
/// 2. It's probably more performant.
|
/// 2. It's probably more performant.
|
||||||
const fn is_protocol(self) -> bool {
|
const fn is_protocol(self) -> bool {
|
||||||
match self {
|
match self {
|
||||||
Self::SupportsIndex | Self::Iterable => true,
|
Self::SupportsIndex | Self::Iterable | Self::Iterator => true,
|
||||||
|
|
||||||
Self::Any
|
Self::Any
|
||||||
| Self::Bool
|
| Self::Bool
|
||||||
|
|
@ -2968,6 +2972,7 @@ impl KnownClass {
|
||||||
Self::ABCMeta => "ABCMeta",
|
Self::ABCMeta => "ABCMeta",
|
||||||
Self::Super => "super",
|
Self::Super => "super",
|
||||||
Self::Iterable => "Iterable",
|
Self::Iterable => "Iterable",
|
||||||
|
Self::Iterator => "Iterator",
|
||||||
// For example, `typing.List` is defined as `List = _Alias()` in typeshed
|
// For example, `typing.List` is defined as `List = _Alias()` in typeshed
|
||||||
Self::StdlibAlias => "_Alias",
|
Self::StdlibAlias => "_Alias",
|
||||||
// This is the name the type of `sys.version_info` has in typeshed,
|
// This is the name the type of `sys.version_info` has in typeshed,
|
||||||
|
|
@ -3203,6 +3208,7 @@ impl KnownClass {
|
||||||
| Self::NamedTuple
|
| Self::NamedTuple
|
||||||
| Self::StdlibAlias
|
| Self::StdlibAlias
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
| Self::SupportsIndex => KnownModule::Typing,
|
| Self::SupportsIndex => KnownModule::Typing,
|
||||||
Self::TypeAliasType
|
Self::TypeAliasType
|
||||||
| Self::TypeVarTuple
|
| Self::TypeVarTuple
|
||||||
|
|
@ -3311,6 +3317,7 @@ impl KnownClass {
|
||||||
| Self::Field
|
| Self::Field
|
||||||
| Self::KwOnly
|
| Self::KwOnly
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
| Self::NamedTupleFallback => false,
|
| Self::NamedTupleFallback => false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -3384,6 +3391,7 @@ impl KnownClass {
|
||||||
| Self::Field
|
| Self::Field
|
||||||
| Self::KwOnly
|
| Self::KwOnly
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
| Self::NamedTupleFallback => false,
|
| Self::NamedTupleFallback => false,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -3435,6 +3443,7 @@ impl KnownClass {
|
||||||
"TypeAliasType" => Self::TypeAliasType,
|
"TypeAliasType" => Self::TypeAliasType,
|
||||||
"TypeVar" => Self::TypeVar,
|
"TypeVar" => Self::TypeVar,
|
||||||
"Iterable" => Self::Iterable,
|
"Iterable" => Self::Iterable,
|
||||||
|
"Iterator" => Self::Iterator,
|
||||||
"ParamSpec" => Self::ParamSpec,
|
"ParamSpec" => Self::ParamSpec,
|
||||||
"ParamSpecArgs" => Self::ParamSpecArgs,
|
"ParamSpecArgs" => Self::ParamSpecArgs,
|
||||||
"ParamSpecKwargs" => Self::ParamSpecKwargs,
|
"ParamSpecKwargs" => Self::ParamSpecKwargs,
|
||||||
|
|
@ -3538,6 +3547,7 @@ impl KnownClass {
|
||||||
| Self::TypeVarTuple
|
| Self::TypeVarTuple
|
||||||
| Self::NamedTuple
|
| Self::NamedTuple
|
||||||
| Self::Iterable
|
| Self::Iterable
|
||||||
|
| Self::Iterator
|
||||||
| Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
|
| Self::NewType => matches!(module, KnownModule::Typing | KnownModule::TypingExtensions),
|
||||||
Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),
|
Self::Deprecated => matches!(module, KnownModule::Warnings | KnownModule::TypingExtensions),
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue