[ty] Detect enums if metaclass is a subtype of EnumType/EnumMeta (#19481)

## Summary

This PR implements the following section from the [typing spec on
enums](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**. Note that `enum.EnumType` was named `enum.EnumMeta` prior
to Python 3.11.

part of https://github.com/astral-sh/ty/issues/183

## Test Plan

New Markdown tests
This commit is contained in:
David Peter 2025-07-23 08:46:51 +02:00 committed by GitHub
parent ba070bb6d5
commit 385d6fa608
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 104 additions and 4 deletions

View file

@ -561,7 +561,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

View file

@ -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

View file

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