mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 13:24:57 +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:
justin
GitHub
parent
c2fa449954
commit
9f0b942b9e
7 changed files with 132 additions and 28 deletions
|
|
@ -2415,7 +2415,7 @@ class Answer(enum.Enum):
|
|||
YES = 1
|
||||
|
||||
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]
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
|
||||
```py
|
||||
from enum import Enum
|
||||
from typing import Literal
|
||||
|
||||
class Color(Enum):
|
||||
RED = 1
|
||||
|
|
@ -11,8 +12,8 @@ class Color(Enum):
|
|||
BLUE = 3
|
||||
|
||||
reveal_type(Color.RED) # revealed: Literal[Color.RED]
|
||||
# TODO: This could be `Literal[1]`
|
||||
reveal_type(Color.RED.value) # revealed: Any
|
||||
reveal_type(Color.RED.name) # revealed: Literal["RED"]
|
||||
reveal_type(Color.RED.value) # revealed: Literal[1]
|
||||
|
||||
# TODO: Should be `Color` or `Literal[Color.RED]`
|
||||
reveal_type(Color["RED"]) # revealed: Unknown
|
||||
|
|
@ -155,6 +156,7 @@ python-version = "3.11"
|
|||
|
||||
```py
|
||||
from enum import Enum, property as enum_property
|
||||
from typing import Any
|
||||
from ty_extensions import enum_members
|
||||
|
||||
class Answer(Enum):
|
||||
|
|
@ -169,6 +171,22 @@ class Answer(Enum):
|
|||
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`
|
||||
|
||||
Attributes defined using `types.DynamicClassAttribute` are not considered members:
|
||||
|
|
@ -481,6 +499,62 @@ callable = Printer.STDERR
|
|||
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
|
||||
|
||||
### 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:
|
||||
# error: [unresolved-attribute]
|
||||
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
|
||||
|
|
|
|||
|
|
@ -4,8 +4,6 @@
|
|||
)]
|
||||
use std::hash::BuildHasherDefault;
|
||||
|
||||
use rustc_hash::FxHasher;
|
||||
|
||||
use crate::lint::{LintRegistry, LintRegistryBuilder};
|
||||
use crate::suppression::{INVALID_IGNORE_COMMENT, UNKNOWN_RULE, UNUSED_IGNORE_COMMENT};
|
||||
pub use db::Db;
|
||||
|
|
@ -19,6 +17,7 @@ pub use program::{
|
|||
PythonVersionWithSource, SearchPathSettings,
|
||||
};
|
||||
pub use python_platform::PythonPlatform;
|
||||
use rustc_hash::FxHasher;
|
||||
pub use semantic_model::{
|
||||
Completion, CompletionKind, HasDefinition, HasType, NameKind, SemanticModel,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -3528,6 +3528,29 @@ impl<'db> Type<'db> {
|
|||
.value_type(db)
|
||||
.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::ProtocolInstance(..)
|
||||
| Type::BooleanLiteral(..)
|
||||
|
|
|
|||
|
|
@ -428,9 +428,8 @@ impl<'db> UnionBuilder<'db> {
|
|||
|
||||
let all_members_are_in_union = metadata
|
||||
.members
|
||||
.difference(&enum_members_in_union)
|
||||
.next()
|
||||
.is_none();
|
||||
.keys()
|
||||
.all(|name| enum_members_in_union.contains(name));
|
||||
|
||||
if all_members_are_in_union {
|
||||
self.add_in_place_impl(
|
||||
|
|
|
|||
|
|
@ -724,7 +724,7 @@ impl<'db> Bindings<'db> {
|
|||
db,
|
||||
metadata
|
||||
.members
|
||||
.iter()
|
||||
.keys()
|
||||
.map(|member| Type::string_literal(db, member)),
|
||||
)
|
||||
} else {
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use ruff_python_ast::name::Name;
|
|||
use rustc_hash::FxHashMap;
|
||||
|
||||
use crate::{
|
||||
Db, FxOrderSet,
|
||||
Db, FxIndexMap,
|
||||
place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
|
||||
semantic_index::{place_table, use_def_map},
|
||||
types::{
|
||||
|
|
@ -11,24 +11,24 @@ use crate::{
|
|||
},
|
||||
};
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub(crate) struct EnumMetadata {
|
||||
pub(crate) members: FxOrderSet<Name>,
|
||||
#[derive(Debug, PartialEq, Eq, salsa::Update)]
|
||||
pub(crate) struct EnumMetadata<'db> {
|
||||
pub(crate) members: FxIndexMap<Name, Type<'db>>,
|
||||
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 {
|
||||
EnumMetadata {
|
||||
members: FxOrderSet::default(),
|
||||
members: FxIndexMap::default(),
|
||||
aliases: FxHashMap::default(),
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
} else {
|
||||
self.aliases.get(name)
|
||||
|
|
@ -36,18 +36,21 @@ impl EnumMetadata {
|
|||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::ref_option)]
|
||||
fn enum_metadata_cycle_recover(
|
||||
_db: &dyn Db,
|
||||
_value: &Option<EnumMetadata>,
|
||||
#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)]
|
||||
fn enum_metadata_cycle_recover<'db>(
|
||||
_db: &'db dyn Db,
|
||||
_value: &Option<EnumMetadata<'db>>,
|
||||
_count: u32,
|
||||
_class: ClassLiteral<'_>,
|
||||
) -> salsa::CycleRecoveryAction<Option<EnumMetadata>> {
|
||||
_class: ClassLiteral<'db>,
|
||||
) -> salsa::CycleRecoveryAction<Option<EnumMetadata<'db>>> {
|
||||
salsa::CycleRecoveryAction::Iterate
|
||||
}
|
||||
|
||||
#[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())
|
||||
}
|
||||
|
||||
|
|
@ -57,7 +60,7 @@ fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option
|
|||
pub(crate) fn enum_metadata<'db>(
|
||||
db: &'db dyn Db,
|
||||
class: ClassLiteral<'db>,
|
||||
) -> Option<EnumMetadata> {
|
||||
) -> Option<EnumMetadata<'db>> {
|
||||
// This is a fast path to avoid traversing the MRO of known classes
|
||||
if class
|
||||
.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() {
|
||||
// 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| {
|
||||
metadata
|
||||
.members
|
||||
.iter()
|
||||
.keys()
|
||||
.filter(move |name| Some(*name) != exclude_member)
|
||||
.map(move |name| Type::EnumLiteral(EnumLiteralType::new(db, class, name.clone())))
|
||||
})
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue