[ty] Detect enums if metaclass is a subtype of EnumType/EnumMeta

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -562,7 +562,83 @@ reveal_type(enum_members(Answer))
 
 ## Custom enum types
 
-To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>
+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
 

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

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;
     }