Skip to content

Commit

Permalink
Use EnumMeta instead of Enum to mark enum classes (#4319)
Browse files Browse the repository at this point in the history
Implement #4311.
  • Loading branch information
elazarg authored and msullivan committed Apr 10, 2018
1 parent 1fd5a30 commit 6f11f35
Show file tree
Hide file tree
Showing 6 changed files with 48 additions and 27 deletions.
15 changes: 0 additions & 15 deletions mypy/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,6 @@ def get_column(self) -> int:
LITERAL_TYPE = 1
LITERAL_NO = 0

# Hard coded name of Enum baseclass.
ENUM_BASECLASS = "enum.Enum"

node_kinds = {
LDEF: 'Ldef',
GDEF: 'Gdef',
Expand Down Expand Up @@ -2152,7 +2149,6 @@ def calculate_mro(self) -> None:
mro = linearize_hierarchy(self)
assert mro, "Could not produce a MRO at all for %s" % (self,)
self.mro = mro
self.is_enum = self._calculate_is_enum()
# The property of falling back to Any is inherited.
self.fallback_to_any = any(baseinfo.fallback_to_any for baseinfo in self.mro)
self.reset_subtype_cache()
Expand All @@ -2178,17 +2174,6 @@ def is_metaclass(self) -> bool:
return (self.has_base('builtins.type') or self.fullname() == 'abc.ABCMeta' or
self.fallback_to_any)

def _calculate_is_enum(self) -> bool:
"""
If this is "enum.Enum" itself, then yes, it's an enum.
If the flag .is_enum has been set on anything in the MRO, it's an enum.
"""
if self.fullname() == ENUM_BASECLASS:
return True
if self.mro:
return any(type_info.is_enum for type_info in self.mro)
return False

def has_base(self, fullname: str) -> bool:
"""Return True if type has a base type with the specified name.
Expand Down
7 changes: 5 additions & 2 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -1085,8 +1085,6 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
# the MRO. Fix MRO if needed.
if info.mro and info.mro[-1].fullname() != 'builtins.object':
info.mro.append(self.object_type().type)
if defn.info.is_enum and defn.type_vars:
self.fail("Enum class cannot be generic", defn)

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

def object_type(self) -> Instance:
return self.named_type('__builtins__.object')
Expand Down
4 changes: 0 additions & 4 deletions mypy/subtypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,10 +183,6 @@ def visit_instance(self, left: Instance) -> bool:
if isinstance(item, AnyType):
return True
if isinstance(item, Instance):
# Special-case enum since we don't have better way of expressing it
if (is_named_instance(left, 'enum.EnumMeta')
and is_named_instance(item, 'enum.Enum')):
return True
return is_named_instance(item, 'builtins.object')
if isinstance(right, CallableType):
# Special case: Instance can be a subtype of Callable.
Expand Down
38 changes: 35 additions & 3 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,42 @@ class Medal(Enum):
gold = 1
silver = 2
bronze = 3
reveal_type(Medal.bronze) # E: Revealed type is '__main__.Medal'
m = Medal.gold
m = 1
[out]
main:7: error: Incompatible types in assignment (expression has type "int", variable has type "Medal")
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")

[case testEnumFromEnumMetaBasics]
from enum import EnumMeta
class Medal(metaclass=EnumMeta):
gold = 1
silver = "hello"
bronze = None
# Without __init__ the definition fails at runtime, but we want to verify that mypy
# uses `enum.EnumMeta` and not `enum.Enum` as the definition of what is enum.
def __init__(self, *args): pass
reveal_type(Medal.bronze) # E: Revealed type is '__main__.Medal'
m = Medal.gold
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")

[case testEnumFromEnumMetaSubclass]
from enum import EnumMeta
class Achievement(metaclass=EnumMeta): pass
class Medal(Achievement):
gold = 1
silver = "hello"
bronze = None
# See comment in testEnumFromEnumMetaBasics
def __init__(self, *args): pass
reveal_type(Medal.bronze) # E: Revealed type is '__main__.Medal'
m = Medal.gold
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")

[case testEnumFromEnumMetaGeneric]
from enum import EnumMeta
from typing import Generic, TypeVar
T = TypeVar("T")
class Medal(Generic[T], metaclass=EnumMeta): # E: Enum class cannot be generic
q = None

[case testEnumNameAndValue]
from enum import Enum
Expand Down
5 changes: 4 additions & 1 deletion test-data/unit/lib-stub/enum.pyi
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
from typing import Any, TypeVar, Union

class Enum:
class EnumMeta(type):
pass

class Enum(metaclass=EnumMeta):
def __new__(cls, value: Any) -> None: pass
def __repr__(self) -> str: pass
def __str__(self) -> str: pass
Expand Down
6 changes: 4 additions & 2 deletions test-data/unit/merge.test
Original file line number Diff line number Diff line change
Expand Up @@ -1440,12 +1440,14 @@ TypeInfo<0>(
Bases(enum.Enum<1>)
Mro(target.A<0>, enum.Enum<1>, builtins.object<2>)
Names(
X<3> (builtins.int<4>)))
X<3> (builtins.int<4>))
MetaclassType(enum.EnumMeta<5>))
==>
TypeInfo<0>(
Name(target.A)
Bases(enum.Enum<1>)
Mro(target.A<0>, enum.Enum<1>, builtins.object<2>)
Names(
X<3> (builtins.int<4>)
Y<5> (builtins.int<4>)))
Y<6> (builtins.int<4>))
MetaclassType(enum.EnumMeta<5>))

0 comments on commit 6f11f35

Please sign in to comment.