Skip to content

Commit 58d94f3

Browse files
committed
use EnumMeta to mark enum classes
1 parent 2c9940b commit 58d94f3

File tree

5 files changed

+21
-25
lines changed

5 files changed

+21
-25
lines changed

mypy/nodes.py

-15
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,6 @@ def get_column(self) -> int:
8181
LITERAL_TYPE = 1
8282
LITERAL_NO = 0
8383

84-
# Hard coded name of Enum baseclass.
85-
ENUM_BASECLASS = "enum.Enum"
86-
8784
node_kinds = {
8885
LDEF: 'Ldef',
8986
GDEF: 'Gdef',
@@ -2085,7 +2082,6 @@ def calculate_mro(self) -> None:
20852082
mro = linearize_hierarchy(self)
20862083
assert mro, "Could not produce a MRO at all for %s" % (self,)
20872084
self.mro = mro
2088-
self.is_enum = self._calculate_is_enum()
20892085

20902086
def calculate_metaclass_type(self) -> 'Optional[mypy.types.Instance]':
20912087
declared = self.declared_metaclass
@@ -2108,17 +2104,6 @@ def is_metaclass(self) -> bool:
21082104
return (self.has_base('builtins.type') or self.fullname() == 'abc.ABCMeta' or
21092105
self.fallback_to_any)
21102106

2111-
def _calculate_is_enum(self) -> bool:
2112-
"""
2113-
If this is "enum.Enum" itself, then yes, it's an enum.
2114-
If the flag .is_enum has been set on anything in the MRO, it's an enum.
2115-
"""
2116-
if self.fullname() == ENUM_BASECLASS:
2117-
return True
2118-
if self.mro:
2119-
return any(type_info.is_enum for type_info in self.mro)
2120-
return False
2121-
21222107
def has_base(self, fullname: str) -> bool:
21232108
"""Return True if type has a base type with the specified name.
21242109

mypy/semanal.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -1083,8 +1083,6 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
10831083
# the MRO. Fix MRO if needed.
10841084
if info.mro and info.mro[-1].fullname() != 'builtins.object':
10851085
info.mro.append(self.object_type().type)
1086-
if defn.info.is_enum and defn.type_vars:
1087-
self.fail("Enum class cannot be generic", defn)
10881086

10891087
def update_metaclass(self, defn: ClassDef) -> None:
10901088
"""Lookup for special metaclass declarations, and update defn fields accordingly.
@@ -1221,6 +1219,11 @@ def analyze_metaclass(self, defn: ClassDef) -> None:
12211219
# do not declare explicit metaclass, but it's harder to catch at this stage
12221220
if defn.metaclass is not None:
12231221
self.fail("Inconsistent metaclass structure for '%s'" % defn.name, defn)
1222+
else:
1223+
if defn.info.metaclass_type.type.fullname() == 'enum.EnumMeta':
1224+
defn.info.is_enum = True
1225+
if defn.type_vars:
1226+
self.fail("Enum class cannot be generic", defn)
12241227

12251228
def object_type(self) -> Instance:
12261229
return self.named_type('__builtins__.object')

mypy/subtypes.py

-4
Original file line numberDiff line numberDiff line change
@@ -183,10 +183,6 @@ def visit_instance(self, left: Instance) -> bool:
183183
if isinstance(item, AnyType):
184184
return True
185185
if isinstance(item, Instance):
186-
# Special-case enum since we don't have better way of expressing it
187-
if (is_named_instance(left, 'enum.EnumMeta')
188-
and is_named_instance(item, 'enum.Enum')):
189-
return True
190186
return is_named_instance(item, 'builtins.object')
191187
if isinstance(right, CallableType):
192188
# Special case: Instance can be a subtype of Callable.

test-data/unit/check-enum.test

+12-3
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,19 @@ class Medal(Enum):
66
gold = 1
77
silver = 2
88
bronze = 3
9+
reveal_type(Medal.bronze) # E: Revealed type is '__main__.Medal'
910
m = Medal.gold
10-
m = 1
11-
[out]
12-
main:7: error: Incompatible types in assignment (expression has type "int", variable has type "Medal")
11+
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")
12+
13+
[case testEnumFromEnumMetaBasics]
14+
from enum import EnumMeta
15+
class Medal(metaclass=EnumMeta):
16+
gold = 1
17+
silver = "hello"
18+
bronze = None
19+
reveal_type(Medal.bronze) # E: Revealed type is '__main__.Medal'
20+
m = Medal.gold
21+
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")
1322

1423
[case testEnumNameAndValue]
1524
from enum import Enum

test-data/unit/lib-stub/enum.pyi

+4-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
from typing import Any, TypeVar, Union
22

3-
class Enum:
3+
class EnumMeta(type):
4+
pass
5+
6+
class Enum(metaclass=EnumMeta):
47
def __new__(cls, value: Any) -> None: pass
58
def __repr__(self) -> str: pass
69
def __str__(self) -> str: pass

0 commit comments

Comments
 (0)