mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-30 13:51:16 +00:00
[ty] infer name
and value
for enum members (#20311)
## summary
- this pr implements the following attributes for `Enum` members:
- `name`
- `_name_`
- `value`
- `_value_`
- adds a TODO test for `my_enum_class_instance.name`
- only implements if the instance is a subclass of `Enum` re: this
[comment](https://github.com/astral-sh/ruff/pull/19481#issuecomment-3103460307)
and existing
[test](c34449ed7c/crates/ty_python_semantic/resources/mdtest/enums.md (L625)
)
### pointers
- https://github.com/astral-sh/ty/issues/876
- https://typing.python.org/en/latest/spec/enums.html#enum-definition
- https://github.com/astral-sh/ruff/pull/19481#issuecomment-3103460307
## test plan
- mdtests
- triaged conformance diffs here:
https://diffswarm.dev/d-01k531ag4nee3xmdeq4f3j66pb
- triaged mypy primer diffs here for django-stubs:
https://diffswarm.dev/d-01k5331n13k9yx8tvnxnkeawp3
- added a TODO test for overriding `.value`
- discord diff seems reasonable
---------
Co-authored-by: David Peter <mail@david-peter.de>
This commit is contained in:
parent
c2fa449954
commit
9f0b942b9e
7 changed files with 132 additions and 28 deletions
|
@ -2415,7 +2415,7 @@ class Answer(enum.Enum):
|
||||||
YES = 1
|
YES = 1
|
||||||
|
|
||||||
reveal_type(Answer.NO) # revealed: Literal[Answer.NO]
|
reveal_type(Answer.NO) # revealed: Literal[Answer.NO]
|
||||||
reveal_type(Answer.NO.value) # revealed: Any
|
reveal_type(Answer.NO.value) # revealed: Literal[0]
|
||||||
reveal_type(Answer.__members__) # revealed: MappingProxyType[str, Unknown]
|
reveal_type(Answer.__members__) # revealed: MappingProxyType[str, Unknown]
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
|
@ -4,6 +4,7 @@
|
||||||
|
|
||||||
```py
|
```py
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
class Color(Enum):
|
class Color(Enum):
|
||||||
RED = 1
|
RED = 1
|
||||||
|
@ -11,8 +12,8 @@ class Color(Enum):
|
||||||
BLUE = 3
|
BLUE = 3
|
||||||
|
|
||||||
reveal_type(Color.RED) # revealed: Literal[Color.RED]
|
reveal_type(Color.RED) # revealed: Literal[Color.RED]
|
||||||
# TODO: This could be `Literal[1]`
|
reveal_type(Color.RED.name) # revealed: Literal["RED"]
|
||||||
reveal_type(Color.RED.value) # revealed: Any
|
reveal_type(Color.RED.value) # revealed: Literal[1]
|
||||||
|
|
||||||
# TODO: Should be `Color` or `Literal[Color.RED]`
|
# TODO: Should be `Color` or `Literal[Color.RED]`
|
||||||
reveal_type(Color["RED"]) # revealed: Unknown
|
reveal_type(Color["RED"]) # revealed: Unknown
|
||||||
|
@ -155,6 +156,7 @@ python-version = "3.11"
|
||||||
|
|
||||||
```py
|
```py
|
||||||
from enum import Enum, property as enum_property
|
from enum import Enum, property as enum_property
|
||||||
|
from typing import Any
|
||||||
from ty_extensions import enum_members
|
from ty_extensions import enum_members
|
||||||
|
|
||||||
class Answer(Enum):
|
class Answer(Enum):
|
||||||
|
@ -169,6 +171,22 @@ class Answer(Enum):
|
||||||
reveal_type(enum_members(Answer))
|
reveal_type(enum_members(Answer))
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Enum attributes defined using `enum.property` take precedence over generated attributes.
|
||||||
|
|
||||||
|
```py
|
||||||
|
from enum import Enum, property as enum_property
|
||||||
|
|
||||||
|
class Choices(Enum):
|
||||||
|
A = 1
|
||||||
|
B = 2
|
||||||
|
|
||||||
|
@enum_property
|
||||||
|
def value(self) -> Any: ...
|
||||||
|
|
||||||
|
# TODO: This should be `Any` - overridden by `@enum_property`
|
||||||
|
reveal_type(Choices.A.value) # revealed: Literal[1]
|
||||||
|
```
|
||||||
|
|
||||||
### `types.DynamicClassAttribute`
|
### `types.DynamicClassAttribute`
|
||||||
|
|
||||||
Attributes defined using `types.DynamicClassAttribute` are not considered members:
|
Attributes defined using `types.DynamicClassAttribute` are not considered members:
|
||||||
|
@ -481,6 +499,62 @@ callable = Printer.STDERR
|
||||||
callable("Another error!")
|
callable("Another error!")
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## Special attributes on enum members
|
||||||
|
|
||||||
|
### `name` and `_name_`
|
||||||
|
|
||||||
|
```py
|
||||||
|
from enum import Enum
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
class Color(Enum):
|
||||||
|
RED = 1
|
||||||
|
GREEN = 2
|
||||||
|
BLUE = 3
|
||||||
|
|
||||||
|
reveal_type(Color.RED._name_) # revealed: Literal["RED"]
|
||||||
|
|
||||||
|
def _(red_or_blue: Literal[Color.RED, Color.BLUE]):
|
||||||
|
reveal_type(red_or_blue.name) # revealed: Literal["RED", "BLUE"]
|
||||||
|
|
||||||
|
def _(any_color: Color):
|
||||||
|
# TODO: Literal["RED", "GREEN", "BLUE"]
|
||||||
|
reveal_type(any_color.name) # revealed: Any
|
||||||
|
```
|
||||||
|
|
||||||
|
### `value` and `_value_`
|
||||||
|
|
||||||
|
```toml
|
||||||
|
[environment]
|
||||||
|
python-version = "3.11"
|
||||||
|
```
|
||||||
|
|
||||||
|
```py
|
||||||
|
from enum import Enum, StrEnum
|
||||||
|
from typing import Literal
|
||||||
|
|
||||||
|
class Color(Enum):
|
||||||
|
RED = 1
|
||||||
|
GREEN = 2
|
||||||
|
BLUE = 3
|
||||||
|
|
||||||
|
reveal_type(Color.RED.value) # revealed: Literal[1]
|
||||||
|
reveal_type(Color.RED._value_) # revealed: Literal[1]
|
||||||
|
|
||||||
|
reveal_type(Color.GREEN.value) # revealed: Literal[2]
|
||||||
|
reveal_type(Color.GREEN._value_) # revealed: Literal[2]
|
||||||
|
|
||||||
|
class Answer(StrEnum):
|
||||||
|
YES = "yes"
|
||||||
|
NO = "no"
|
||||||
|
|
||||||
|
reveal_type(Answer.YES.value) # revealed: Literal["yes"]
|
||||||
|
reveal_type(Answer.YES._value_) # revealed: Literal["yes"]
|
||||||
|
|
||||||
|
reveal_type(Answer.NO.value) # revealed: Literal["no"]
|
||||||
|
reveal_type(Answer.NO._value_) # revealed: Literal["no"]
|
||||||
|
```
|
||||||
|
|
||||||
## Properties of enum types
|
## Properties of enum types
|
||||||
|
|
||||||
### Implicitly final
|
### Implicitly final
|
||||||
|
@ -609,6 +683,12 @@ reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWit
|
||||||
# Attributes like `.value` can *not* be accessed on members of these enums:
|
# Attributes like `.value` can *not* be accessed on members of these enums:
|
||||||
# error: [unresolved-attribute]
|
# error: [unresolved-attribute]
|
||||||
EnumWithSubclassOfEnumMetaMetaclass.NO.value
|
EnumWithSubclassOfEnumMetaMetaclass.NO.value
|
||||||
|
# error: [unresolved-attribute]
|
||||||
|
EnumWithSubclassOfEnumMetaMetaclass.NO._value_
|
||||||
|
# error: [unresolved-attribute]
|
||||||
|
EnumWithSubclassOfEnumMetaMetaclass.NO.name
|
||||||
|
# error: [unresolved-attribute]
|
||||||
|
EnumWithSubclassOfEnumMetaMetaclass.NO._name_
|
||||||
```
|
```
|
||||||
|
|
||||||
### Enums with (subclasses of) `EnumType` as metaclass
|
### Enums with (subclasses of) `EnumType` as metaclass
|
||||||
|
|
|
@ -4,8 +4,6 @@
|
||||||
)]
|
)]
|
||||||
use std::hash::BuildHasherDefault;
|
use std::hash::BuildHasherDefault;
|
||||||
|
|
||||||
use rustc_hash::FxHasher;
|
|
||||||
|
|
||||||
use crate::lint::{LintRegistry, LintRegistryBuilder};
|
use crate::lint::{LintRegistry, LintRegistryBuilder};
|
||||||
use crate::suppression::{INVALID_IGNORE_COMMENT, UNKNOWN_RULE, UNUSED_IGNORE_COMMENT};
|
use crate::suppression::{INVALID_IGNORE_COMMENT, UNKNOWN_RULE, UNUSED_IGNORE_COMMENT};
|
||||||
pub use db::Db;
|
pub use db::Db;
|
||||||
|
@ -19,6 +17,7 @@ pub use program::{
|
||||||
PythonVersionWithSource, SearchPathSettings,
|
PythonVersionWithSource, SearchPathSettings,
|
||||||
};
|
};
|
||||||
pub use python_platform::PythonPlatform;
|
pub use python_platform::PythonPlatform;
|
||||||
|
use rustc_hash::FxHasher;
|
||||||
pub use semantic_model::{
|
pub use semantic_model::{
|
||||||
Completion, CompletionKind, HasDefinition, HasType, NameKind, SemanticModel,
|
Completion, CompletionKind, HasDefinition, HasType, NameKind, SemanticModel,
|
||||||
};
|
};
|
||||||
|
|
|
@ -3528,6 +3528,29 @@ impl<'db> Type<'db> {
|
||||||
.value_type(db)
|
.value_type(db)
|
||||||
.member_lookup_with_policy(db, name, policy),
|
.member_lookup_with_policy(db, name, policy),
|
||||||
|
|
||||||
|
Type::EnumLiteral(enum_literal)
|
||||||
|
if matches!(name_str, "name" | "_name_")
|
||||||
|
&& Type::ClassLiteral(enum_literal.enum_class(db))
|
||||||
|
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) =>
|
||||||
|
{
|
||||||
|
Place::bound(Type::StringLiteral(StringLiteralType::new(
|
||||||
|
db,
|
||||||
|
enum_literal.name(db).as_str(),
|
||||||
|
)))
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
|
||||||
|
Type::EnumLiteral(enum_literal)
|
||||||
|
if matches!(name_str, "value" | "_value_")
|
||||||
|
&& Type::ClassLiteral(enum_literal.enum_class(db))
|
||||||
|
.is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) =>
|
||||||
|
{
|
||||||
|
enum_metadata(db, enum_literal.enum_class(db))
|
||||||
|
.and_then(|metadata| metadata.members.get(enum_literal.name(db)))
|
||||||
|
.map_or_else(|| Place::Unbound, Place::bound)
|
||||||
|
.into()
|
||||||
|
}
|
||||||
|
|
||||||
Type::NominalInstance(..)
|
Type::NominalInstance(..)
|
||||||
| Type::ProtocolInstance(..)
|
| Type::ProtocolInstance(..)
|
||||||
| Type::BooleanLiteral(..)
|
| Type::BooleanLiteral(..)
|
||||||
|
|
|
@ -428,9 +428,8 @@ impl<'db> UnionBuilder<'db> {
|
||||||
|
|
||||||
let all_members_are_in_union = metadata
|
let all_members_are_in_union = metadata
|
||||||
.members
|
.members
|
||||||
.difference(&enum_members_in_union)
|
.keys()
|
||||||
.next()
|
.all(|name| enum_members_in_union.contains(name));
|
||||||
.is_none();
|
|
||||||
|
|
||||||
if all_members_are_in_union {
|
if all_members_are_in_union {
|
||||||
self.add_in_place_impl(
|
self.add_in_place_impl(
|
||||||
|
|
|
@ -724,7 +724,7 @@ impl<'db> Bindings<'db> {
|
||||||
db,
|
db,
|
||||||
metadata
|
metadata
|
||||||
.members
|
.members
|
||||||
.iter()
|
.keys()
|
||||||
.map(|member| Type::string_literal(db, member)),
|
.map(|member| Type::string_literal(db, member)),
|
||||||
)
|
)
|
||||||
} else {
|
} else {
|
||||||
|
|
|
@ -2,7 +2,7 @@ use ruff_python_ast::name::Name;
|
||||||
use rustc_hash::FxHashMap;
|
use rustc_hash::FxHashMap;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
Db, FxOrderSet,
|
Db, FxIndexMap,
|
||||||
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
|
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
|
||||||
semantic_index::{place_table, use_def_map},
|
semantic_index::{place_table, use_def_map},
|
||||||
types::{
|
types::{
|
||||||
|
@ -11,24 +11,24 @@ use crate::{
|
||||||
},
|
},
|
||||||
};
|
};
|
||||||
|
|
||||||
#[derive(Debug, PartialEq, Eq)]
|
#[derive(Debug, PartialEq, Eq, salsa::Update)]
|
||||||
pub(crate) struct EnumMetadata {
|
pub(crate) struct EnumMetadata<'db> {
|
||||||
pub(crate) members: FxOrderSet<Name>,
|
pub(crate) members: FxIndexMap<Name, Type<'db>>,
|
||||||
pub(crate) aliases: FxHashMap<Name, Name>,
|
pub(crate) aliases: FxHashMap<Name, Name>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl get_size2::GetSize for EnumMetadata {}
|
impl get_size2::GetSize for EnumMetadata<'_> {}
|
||||||
|
|
||||||
impl EnumMetadata {
|
impl EnumMetadata<'_> {
|
||||||
fn empty() -> Self {
|
fn empty() -> Self {
|
||||||
EnumMetadata {
|
EnumMetadata {
|
||||||
members: FxOrderSet::default(),
|
members: FxIndexMap::default(),
|
||||||
aliases: FxHashMap::default(),
|
aliases: FxHashMap::default(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
|
pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> {
|
||||||
if self.members.contains(name) {
|
if self.members.contains_key(name) {
|
||||||
Some(name)
|
Some(name)
|
||||||
} else {
|
} else {
|
||||||
self.aliases.get(name)
|
self.aliases.get(name)
|
||||||
|
@ -36,18 +36,21 @@ impl EnumMetadata {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(clippy::ref_option)]
|
#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)]
|
||||||
fn enum_metadata_cycle_recover(
|
fn enum_metadata_cycle_recover<'db>(
|
||||||
_db: &dyn Db,
|
_db: &'db dyn Db,
|
||||||
_value: &Option<EnumMetadata>,
|
_value: &Option<EnumMetadata<'db>>,
|
||||||
_count: u32,
|
_count: u32,
|
||||||
_class: ClassLiteral<'_>,
|
_class: ClassLiteral<'db>,
|
||||||
) -> salsa::CycleRecoveryAction<Option<EnumMetadata>> {
|
) -> salsa::CycleRecoveryAction<Option<EnumMetadata<'db>>> {
|
||||||
salsa::CycleRecoveryAction::Iterate
|
salsa::CycleRecoveryAction::Iterate
|
||||||
}
|
}
|
||||||
|
|
||||||
#[allow(clippy::unnecessary_wraps)]
|
#[allow(clippy::unnecessary_wraps)]
|
||||||
fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option<EnumMetadata> {
|
fn enum_metadata_cycle_initial<'db>(
|
||||||
|
_db: &'db dyn Db,
|
||||||
|
_class: ClassLiteral<'db>,
|
||||||
|
) -> Option<EnumMetadata<'db>> {
|
||||||
Some(EnumMetadata::empty())
|
Some(EnumMetadata::empty())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -57,7 +60,7 @@ fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option
|
||||||
pub(crate) fn enum_metadata<'db>(
|
pub(crate) fn enum_metadata<'db>(
|
||||||
db: &'db dyn Db,
|
db: &'db dyn Db,
|
||||||
class: ClassLiteral<'db>,
|
class: ClassLiteral<'db>,
|
||||||
) -> Option<EnumMetadata> {
|
) -> Option<EnumMetadata<'db>> {
|
||||||
// This is a fast path to avoid traversing the MRO of known classes
|
// This is a fast path to avoid traversing the MRO of known classes
|
||||||
if class
|
if class
|
||||||
.known(db)
|
.known(db)
|
||||||
|
@ -217,9 +220,9 @@ pub(crate) fn enum_metadata<'db>(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
Some(name.clone())
|
Some((name.clone(), value_ty))
|
||||||
})
|
})
|
||||||
.collect::<FxOrderSet<_>>();
|
.collect::<FxIndexMap<_, _>>();
|
||||||
|
|
||||||
if members.is_empty() {
|
if members.is_empty() {
|
||||||
// Enum subclasses without members are not considered enums.
|
// Enum subclasses without members are not considered enums.
|
||||||
|
@ -237,7 +240,7 @@ pub(crate) fn enum_member_literals<'a, 'db: 'a>(
|
||||||
enum_metadata(db, class).map(|metadata| {
|
enum_metadata(db, class).map(|metadata| {
|
||||||
metadata
|
metadata
|
||||||
.members
|
.members
|
||||||
.iter()
|
.keys()
|
||||||
.filter(move |name| Some(*name) != exclude_member)
|
.filter(move |name| Some(*name) != exclude_member)
|
||||||
.map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone())))
|
.map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone())))
|
||||||
})
|
})
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue