Skip to content

Commit 018c0e8

Browse files
sharkdpAlexWaygood
authored andcommitted
[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 astral-sh/ty#183 ## Test Plan New Markdown tests
1 parent f122273 commit 018c0e8

File tree

3 files changed

+104
-4
lines changed

3 files changed

+104
-4
lines changed

crates/ty_python_semantic/resources/mdtest/enums.md

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -561,7 +561,83 @@ reveal_type(enum_members(Answer))
561561

562562
## Custom enum types
563563

564-
To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>
564+
Enum classes can also be defined using a subclass of `enum.Enum` or any class that uses
565+
`enum.EnumType` (or a subclass thereof) as a metaclass. `enum.EnumType` was called `enum.EnumMeta`
566+
prior to Python 3.11.
567+
568+
### Subclasses of `Enum`
569+
570+
```py
571+
from enum import Enum, EnumMeta
572+
573+
class CustomEnumSubclass(Enum):
574+
def custom_method(self) -> int:
575+
return 0
576+
577+
class EnumWithCustomEnumSubclass(CustomEnumSubclass):
578+
NO = 0
579+
YES = 1
580+
581+
reveal_type(EnumWithCustomEnumSubclass.NO) # revealed: Literal[EnumWithCustomEnumSubclass.NO]
582+
reveal_type(EnumWithCustomEnumSubclass.NO.custom_method()) # revealed: int
583+
```
584+
585+
### Enums with (subclasses of) `EnumMeta` as metaclass
586+
587+
```toml
588+
[environment]
589+
python-version = "3.9"
590+
```
591+
592+
```py
593+
from enum import Enum, EnumMeta
594+
595+
class EnumWithEnumMetaMetaclass(metaclass=EnumMeta):
596+
NO = 0
597+
YES = 1
598+
599+
reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaMetaclass.NO]
600+
601+
class SubclassOfEnumMeta(EnumMeta): ...
602+
603+
class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta):
604+
NO = 0
605+
YES = 1
606+
607+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO]
608+
609+
# Attributes like `.value` can *not* be accessed on members of these enums:
610+
# error: [unresolved-attribute]
611+
EnumWithSubclassOfEnumMetaMetaclass.NO.value
612+
```
613+
614+
### Enums with (subclasses of) `EnumType` as metaclass
615+
616+
```toml
617+
[environment]
618+
python-version = "3.11"
619+
```
620+
621+
```py
622+
from enum import Enum, EnumType
623+
624+
class EnumWithEnumMetaMetaclass(metaclass=EnumType):
625+
NO = 0
626+
YES = 1
627+
628+
reveal_type(EnumWithEnumMetaMetaclass.NO) # revealed: Literal[EnumWithEnumMetaMetaclass.NO]
629+
630+
class SubclassOfEnumMeta(EnumType): ...
631+
632+
class EnumWithSubclassOfEnumMetaMetaclass(metaclass=SubclassOfEnumMeta):
633+
NO = 0
634+
YES = 1
635+
636+
reveal_type(EnumWithSubclassOfEnumMetaMetaclass.NO) # revealed: Literal[EnumWithSubclassOfEnumMetaMetaclass.NO]
637+
638+
# error: [unresolved-attribute]
639+
EnumWithSubclassOfEnumMetaMetaclass.NO.value
640+
```
565641

566642
## Function syntax
567643

crates/ty_python_semantic/src/types/class.rs

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2656,6 +2656,7 @@ pub enum KnownClass {
26562656
Super,
26572657
// enum
26582658
Enum,
2659+
EnumType,
26592660
Auto,
26602661
Member,
26612662
Nonmember,
@@ -2781,6 +2782,7 @@ impl KnownClass {
27812782
| Self::Deque
27822783
| Self::Float
27832784
| Self::Enum
2785+
| Self::EnumType
27842786
| Self::Auto
27852787
| Self::Member
27862788
| Self::Nonmember
@@ -2865,6 +2867,7 @@ impl KnownClass {
28652867
Self::ABCMeta
28662868
| Self::Any
28672869
| Self::Enum
2870+
| Self::EnumType
28682871
| Self::Auto
28692872
| Self::Member
28702873
| Self::Nonmember
@@ -2914,6 +2917,7 @@ impl KnownClass {
29142917
| KnownClass::Deprecated
29152918
| KnownClass::Super
29162919
| KnownClass::Enum
2920+
| KnownClass::EnumType
29172921
| KnownClass::Auto
29182922
| KnownClass::Member
29192923
| KnownClass::Nonmember
@@ -3022,6 +3026,7 @@ impl KnownClass {
30223026
| Self::Deque
30233027
| Self::OrderedDict
30243028
| Self::Enum
3029+
| Self::EnumType
30253030
| Self::Auto
30263031
| Self::Member
30273032
| Self::Nonmember
@@ -3091,6 +3096,13 @@ impl KnownClass {
30913096
Self::Deque => "deque",
30923097
Self::OrderedDict => "OrderedDict",
30933098
Self::Enum => "Enum",
3099+
Self::EnumType => {
3100+
if Program::get(db).python_version(db) >= PythonVersion::PY311 {
3101+
"EnumType"
3102+
} else {
3103+
"EnumMeta"
3104+
}
3105+
}
30943106
Self::Auto => "auto",
30953107
Self::Member => "member",
30963108
Self::Nonmember => "nonmember",
@@ -3316,7 +3328,9 @@ impl KnownClass {
33163328
| Self::Property => KnownModule::Builtins,
33173329
Self::VersionInfo => KnownModule::Sys,
33183330
Self::ABCMeta => KnownModule::Abc,
3319-
Self::Enum | Self::Auto | Self::Member | Self::Nonmember => KnownModule::Enum,
3331+
Self::Enum | Self::EnumType | Self::Auto | Self::Member | Self::Nonmember => {
3332+
KnownModule::Enum
3333+
}
33203334
Self::GenericAlias
33213335
| Self::ModuleType
33223336
| Self::FunctionType
@@ -3432,6 +3446,7 @@ impl KnownClass {
34323446
| Self::ParamSpecKwargs
34333447
| Self::TypeVarTuple
34343448
| Self::Enum
3449+
| Self::EnumType
34353450
| Self::Auto
34363451
| Self::Member
34373452
| Self::Nonmember
@@ -3505,6 +3520,7 @@ impl KnownClass {
35053520
| Self::ParamSpecKwargs
35063521
| Self::TypeVarTuple
35073522
| Self::Enum
3523+
| Self::EnumType
35083524
| Self::Auto
35093525
| Self::Member
35103526
| Self::Nonmember
@@ -3583,6 +3599,10 @@ impl KnownClass {
35833599
"_NoDefaultType" => Self::NoDefaultType,
35843600
"SupportsIndex" => Self::SupportsIndex,
35853601
"Enum" => Self::Enum,
3602+
"EnumMeta" => Self::EnumType,
3603+
"EnumType" if Program::get(db).python_version(db) >= PythonVersion::PY311 => {
3604+
Self::EnumType
3605+
}
35863606
"auto" => Self::Auto,
35873607
"member" => Self::Member,
35883608
"nonmember" => Self::Nonmember,
@@ -3647,6 +3667,7 @@ impl KnownClass {
36473667
| Self::MethodType
36483668
| Self::MethodWrapperType
36493669
| Self::Enum
3670+
| Self::EnumType
36503671
| Self::Auto
36513672
| Self::Member
36523673
| Self::Nonmember

crates/ty_python_semantic/src/types/enums.rs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,11 @@ pub(crate) fn enum_metadata<'db>(
6666
return None;
6767
}
6868

69-
// TODO: This check needs to be extended (`EnumMeta`/`EnumType`)
70-
if !Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db)) {
69+
if !Type::ClassLiteral(class).is_subtype_of(db, KnownClass::Enum.to_subclass_of(db))
70+
&& !class
71+
.metaclass(db)
72+
.is_subtype_of(db, KnownClass::EnumType.to_subclass_of(db))
73+
{
7174
return None;
7275
}
7376

0 commit comments

Comments
 (0)