diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 482fdd49609509..b266f8d007720f 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -562,7 +562,83 @@ reveal_type(enum_members(Answer)) ## Custom enum types -To do: +Enum classes can also be defined using a subclass of `enum.Enum` or any class that uses +`enum.EnumType` (or a subclass thereof) as a metaclass. `enum.EnumType` was called `enum.EnumMeta` +prior to Python 3.11. + +### Subclasses of `Enum` + +```py +from enum import Enum, EnumMeta + +class CustomEnumSubclass(Enum): + def custom_method(self) -> int: + return 0 + +class EnumWithCustomEnumSubclass(CustomEnumSubclass): + NO = 0 + YES = 1 + +reveal_type(EnumWithCustomEnumSubclass.NO) # revealed: Literal[EnumWithCustomEnumSubclass.NO] +reveal_type(EnumWithCustomEnumSubclass.NO.custom_method()) # revealed: int +``` + +### Enums with (subclasses of) `EnumMeta` as metaclass + +```toml +[environment] +python-version = "3.9" +``` + +```py +from enum import Enum, EnumMeta + +class EnumWithEnumMetaMetaclass(metaclass=EnumMeta): + NO = 0 + YES = 1 + +reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaMetaclass.NO] + +class SubclassOfEnumMeta(EnumMeta): ... + +class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta): + NO = 0 + YES = 1 + +reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO] + +# Attributes like `.value` can *not* be accessed on members of these enums: +# error: [unresolved-attribute] +EnumWithSubclassOfEnumMetaMetaclass.NO.value +``` + +### Enums with (subclasses of) `EnumType` as metaclass + +```toml +[environment] +python-version = "3.11" +``` + +```py +from enum import Enum, EnumType + +class EnumWithEnumMetaMetaclass(metaclass=EnumType): + NO = 0 + YES = 1 + +reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaMetaclass.NO] + +class SubclassOfEnumMeta(EnumType): ... + +class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta): + NO = 0 + YES = 1 + +reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO] + +# error: [unresolved-attribute] +EnumWithSubclassOfEnumMetaMetaclass.NO.value +``` ## Function syntax diff --git a/crates/ty_python_semantic/src/types/class.rs b/crates/ty_python_semantic/src/types/class.rs index f4f21cfab2acdc..c2b2d535488ab0 100644 --- a/crates/ty_python_semantic/src/types/class.rs +++ b/crates/ty_python_semantic/src/types/class.rs @@ -2531,6 +2531,7 @@ pub enum KnownClass { Super, // enum Enum, + EnumType, Auto, Member, Nonmember, @@ -2655,6 +2656,7 @@ impl KnownClass { | Self::Deque | Self::Float | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2738,6 +2740,7 @@ impl KnownClass { Self::ABCMeta | Self::Any | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2786,6 +2789,7 @@ impl KnownClass { | KnownClass::Deprecated | KnownClass::Super | KnownClass::Enum + | KnownClass::EnumType | KnownClass::Auto | KnownClass::Member | KnownClass::Nonmember @@ -2893,6 +2897,7 @@ impl KnownClass { | Self::Deque | Self::OrderedDict | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2962,6 +2967,13 @@ impl KnownClass { Self::Deque => "deque", Self::OrderedDict => "OrderedDict", Self::Enum => "Enum", + Self::EnumType => { + if Program::get(db).python_version(db) >= PythonVersion::PY311 { + "EnumType" + } else { + "EnumMeta" + } + } Self::Auto => "auto", Self::Member => "member", Self::Nonmember => "nonmember", @@ -3186,7 +3198,9 @@ impl KnownClass { | Self::Property => KnownModule::Builtins, Self::VersionInfo => KnownModule::Sys, Self::ABCMeta => KnownModule::Abc, - Self::Enum | Self::Auto | Self::Member | Self::Nonmember => KnownModule::Enum, + Self::Enum | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember => { + KnownModule::Enum + } Self::GenericAlias | Self::ModuleType | Self::FunctionType @@ -3301,6 +3315,7 @@ impl KnownClass { | Self::ParamSpecKwargs | Self::TypeVarTuple | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -3373,6 +3388,7 @@ impl KnownClass { | Self::ParamSpecKwargs | Self::TypeVarTuple | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -3449,6 +3465,10 @@ impl KnownClass { "_NoDefaultType" => Self::NoDefaultType, "SupportsIndex" => Self::SupportsIndex, "Enum" => Self::Enum, + "EnumMeta" => Self::EnumType, + "EnumType" if Program::get(db).python_version(db) >= PythonVersion::PY311 => { + Self::EnumType + } "auto" => Self::Auto, "member" => Self::Member, "nonmember" => Self::Nonmember, @@ -3513,6 +3533,7 @@ impl KnownClass { | Self::MethodType | Self::MethodWrapperType | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 59c814c147995b..d35bd92f57ddc0 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -66,8 +66,11 @@ pub(crate) fn enum_metadata<'db>( return None; } - // TODO: This check needs to be extended (`EnumMeta`/`EnumType`) - if !Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) { + 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)) + { return None; }