Use EnumMeta instead of Enum to mark enum classes (#4319)

Implement #4311.

Author: elazarg

mypy/nodes.py

@@ -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',
@@ -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()
@@ -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.
 

mypy/semanal.py

@@ -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.
@@ -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')

mypy/subtypes.py

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

test-data/unit/check-enum.test

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

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

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

test-data/unit/merge.test

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