Skip to content

Commit

Permalink
Do EnumType as well.
Browse files Browse the repository at this point in the history
  • Loading branch information
AlexWaygood committed Nov 1, 2023
1 parent 959d52c commit 6f76d33
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 3 deletions.
7 changes: 6 additions & 1 deletion pyi.py
Original file line number Diff line number Diff line change
Expand Up @@ -823,7 +823,12 @@ def _is_enum_class(node: ast.ClassDef) -> bool:
return any(_is_enum_base(base) for base in node.bases)


_COMMON_METACLASSES = {"type": "builtins", "ABCMeta": "abc", "EnumMeta": "enum"}
_COMMON_METACLASSES = {
"type": "builtins",
"ABCMeta": "abc",
"EnumMeta": "enum",
"EnumType": "enum",
}


def _is_metaclass_base(node: ast.expr) -> bool:
Expand Down
11 changes: 9 additions & 2 deletions tests/classdefs.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,9 @@
import abc
import builtins
import collections.abc
import enum
import typing
from abc import abstractmethod
from abc import ABCMeta, abstractmethod
from collections.abc import (
AsyncGenerator,
AsyncIterable,
Expand Down Expand Up @@ -139,12 +140,18 @@ class MetaclassInWhichSelfCannotBeUsed2(EnumMeta):
async def __aenter__(self) -> MetaclassInWhichSelfCannotBeUsed2: ...
def __isub__(self, other: MetaclassInWhichSelfCannotBeUsed2) -> MetaclassInWhichSelfCannotBeUsed2: ...

class MetaclassInWhichSelfCannotBeUsed3(abc.ABCMeta):
class MetaclassInWhichSelfCannotBeUsed3(enum.EnumType):
def __new__(cls) -> MetaclassInWhichSelfCannotBeUsed3: ...
def __enter__(self) -> MetaclassInWhichSelfCannotBeUsed3: ...
async def __aenter__(self) -> MetaclassInWhichSelfCannotBeUsed3: ...
def __isub__(self, other: MetaclassInWhichSelfCannotBeUsed3) -> MetaclassInWhichSelfCannotBeUsed3: ...

class MetaclassInWhichSelfCannotBeUsed4(ABCMeta):
def __new__(cls) -> MetaclassInWhichSelfCannotBeUsed4: ...
def __enter__(self) -> MetaclassInWhichSelfCannotBeUsed4: ...
async def __aenter__(self) -> MetaclassInWhichSelfCannotBeUsed4: ...
def __isub__(self, other: MetaclassInWhichSelfCannotBeUsed4) -> MetaclassInWhichSelfCannotBeUsed4: ...

class Abstract(Iterator[str]):
@abstractmethod
def __iter__(self) -> Iterator[str]: ...
Expand Down

0 comments on commit 6f76d33

Please sign in to comment.