diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index 1ef154ee36..0e49bf490d 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -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] ``` diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index a70a92d430..0e314e982c 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -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 diff --git a/crates/ty_python_semantic/src/lib.rs b/crates/ty_python_semantic/src/lib.rs index 48a2cc9eb1..30300eef13 100644 --- a/crates/ty_python_semantic/src/lib.rs +++ b/crates/ty_python_semantic/src/lib.rs @@ -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, }; diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index fdac4d843c..addbc4491b 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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(..) diff --git a/crates/ty_python_semantic/src/types/builder.rs b/crates/ty_python_semantic/src/types/builder.rs index a729191e00..cc1c35790a 100644 --- a/crates/ty_python_semantic/src/types/builder.rs +++ b/crates/ty_python_semantic/src/types/builder.rs @@ -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( diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 9866695233..e7c223a8b4 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/crates/ty_python_semantic/src/types/call/bind.rs @@ -724,7 +724,7 @@ impl<'db> Bindings<'db> { db, metadata .members - .iter() + .keys() .map(|member| Type::string_literal(db, member)), ) } else { diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index c62851fe8d..828ed580b8 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -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, +#[derive(Debug, PartialEq, Eq, salsa::Update)] +pub(crate) struct EnumMetadata<'db> { + pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, } -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, +#[allow(clippy::ref_option, clippy::trivially_copy_pass_by_ref)] +fn enum_metadata_cycle_recover<'db>( + _db: &'db dyn Db, + _value: &Option>, _count: u32, - _class: ClassLiteral<'_>, -) -> salsa::CycleRecoveryAction> { + _class: ClassLiteral<'db>, +) -> salsa::CycleRecoveryAction>> { salsa::CycleRecoveryAction::Iterate } #[allow(clippy::unnecessary_wraps)] -fn enum_metadata_cycle_initial(_db: &dyn Db, _class: ClassLiteral<'_>) -> Option { +fn enum_metadata_cycle_initial<'db>( + _db: &'db dyn Db, + _class: ClassLiteral<'db>, +) -> Option> { 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 { +) -> Option> { // 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::>(); + .collect::>(); 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()))) })