diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 43cd712a1d..8809f41350 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -451,6 +451,10 @@ class Answer(Enum): # revealed: tuple[Literal["YES"], Literal["NO"]] reveal_type(enum_members(Answer)) + +# `nonmember` attributes are unwrapped to the inner value type when accessed. +# revealed: int +reveal_type(Answer.OTHER) ``` `member` can also be used as a decorator: diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index 474c30d25c..dd6b192d2d 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -20,7 +20,9 @@ use crate::types::bound_super::BoundSuperError; use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension}; use crate::types::context::InferContext; use crate::types::diagnostic::{INVALID_TYPE_ALIAS_TYPE, SUPER_CALL_IN_NAMED_TUPLE_METHOD}; -use crate::types::enums::enum_metadata; +use crate::types::enums::{ + enum_metadata, is_enum_class_by_inheritance, try_unwrap_nonmember_value, +}; use crate::types::function::{ DataclassTransformerFlags, DataclassTransformerParams, KnownFunction, }; @@ -2355,6 +2357,18 @@ impl<'db> ClassLiteral<'db> { // The symbol was not found in the class scope. It might still be implicitly defined in `@classmethod`s. return Self::implicit_attribute(db, body_scope, name, MethodDecorator::ClassMethod); } + + // For enum classes, `nonmember(value)` creates a non-member attribute. + // At runtime, the enum metaclass unwraps the value, so accessing the attribute + // returns the inner value, not the `nonmember` wrapper. + if let Some(ty) = member.inner.place.ignore_possibly_undefined() { + if let Some(value_ty) = try_unwrap_nonmember_value(db, ty) { + if is_enum_class_by_inheritance(db, self) { + return Member::definitely_declared(value_ty); + } + } + } + member } diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index fcc8552b52..d478389f7a 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -60,11 +60,7 @@ pub(crate) fn enum_metadata<'db>( return None; } - if !Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) - && !class - .metaclass(db) - .is_subtype_of(db, KnownClass::EnumType.to_subclass_of(db)) - { + if !is_enum_class_by_inheritance(db, class) { return None; } @@ -294,3 +290,63 @@ pub(crate) fn is_enum_class<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool { _ => false, } } + +/// Checks if a class is an enum class by inheritance (either a subtype of `Enum` +/// or has a metaclass that is a subtype of `EnumType`). +/// +/// This is a lighter-weight check than `enum_metadata`, which additionally +/// verifies that the class has members. +pub(crate) fn is_enum_class_by_inheritance<'db>(db: &'db dyn Db, class: ClassLiteral<'db>) -> bool { + Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) + || class + .metaclass(db) + .is_subtype_of(db, KnownClass::EnumType.to_subclass_of(db)) +} + +/// Extracts the inner value type from an `enum.nonmember()` wrapper. +/// +/// At runtime, the enum metaclass unwraps `nonmember(value)`, so accessing the attribute +/// returns the inner value, not the `nonmember` wrapper. +/// +/// Returns `Some(value_type)` if the type is a `nonmember[T]`, otherwise `None`. +pub(crate) fn try_unwrap_nonmember_value<'db>(db: &'db dyn Db, ty: Type<'db>) -> Option> { + match ty { + Type::NominalInstance(instance) if instance.has_known_class(db, KnownClass::Nonmember) => { + Some( + ty.member(db, "value") + .place + .ignore_possibly_undefined() + .unwrap_or(Type::unknown()), + ) + } + Type::Union(union) => { + // TODO: This is a hack. The proper fix is to avoid unioning Unknown from + // declarations into Place when we have concrete bindings. + // + // For now, we filter out Unknown and expect exactly one nonmember type + // to remain. If there are other non-Unknown types mixed in, we bail out. + let mut non_unknown = union.elements(db).iter().filter(|elem| !elem.is_unknown()); + + let first = non_unknown.next()?; + + // Ensure there's exactly one non-Unknown element. + if non_unknown.next().is_some() { + return None; + } + + if let Type::NominalInstance(instance) = first { + if instance.has_known_class(db, KnownClass::Nonmember) { + return Some( + first + .member(db, "value") + .place + .ignore_possibly_undefined() + .unwrap_or(Type::unknown()), + ); + } + } + None + } + _ => None, + } +}