diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index e804e984a8..6981f5c06a 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -561,7 +561,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 34ba5797bd..bc3998c43a 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, @@ -2656,6 +2657,7 @@ impl KnownClass { | Self::Deque | Self::Float | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2740,6 +2742,7 @@ impl KnownClass { Self::ABCMeta | Self::Any | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2789,6 +2792,7 @@ impl KnownClass { | KnownClass::Deprecated | KnownClass::Super | KnownClass::Enum + | KnownClass::EnumType | KnownClass::Auto | KnownClass::Member | KnownClass::Nonmember @@ -2897,6 +2901,7 @@ impl KnownClass { | Self::Deque | Self::OrderedDict | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -2966,6 +2971,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", @@ -3191,7 +3203,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 @@ -3307,6 +3321,7 @@ impl KnownClass { | Self::ParamSpecKwargs | Self::TypeVarTuple | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -3380,6 +3395,7 @@ impl KnownClass { | Self::ParamSpecKwargs | Self::TypeVarTuple | Self::Enum + | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember @@ -3458,6 +3474,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, @@ -3522,6 +3542,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 5c026beff5..c10fcdce39 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; }