Enum: deal with subclasses of Enum or IntEnum (#1601)

* Enum: deal with subclasses of Enum or IntEnum

* better fullname()

* Calculate MRO per-class instead of per-instance. Amend to trigger rebuild

Author: wittekm

mypy/nodes.py

@@ -56,6 +56,9 @@ def get_line(self) -> int: pass
 LITERAL_TYPE = 1
 LITERAL_NO = 0
 
+# Hard coded name of Enum baseclass.
+ENUM_BASECLASS = "enum.Enum"
+
 node_kinds = {
     LDEF: 'Ldef',
     GDEF: 'Gdef',
@@ -1873,6 +1876,18 @@ 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()
+
+    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

@@ -127,12 +127,6 @@
     'builtins.bytearray': 'builtins.str',
 })
 
-# Hard coded list of Enum baseclasses.
-ENUM_BASECLASSES = [
-    'enum.Enum',
-    'enum.IntEnum',
-]
-
 # When analyzing a function, should we analyze the whole function in one go, or
 # should we only perform one phase of the analysis? The latter is used for
 # nested functions. In the first phase we add the function to the symbol table
@@ -765,8 +759,6 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
                 defn.info.fallback_to_any = True
             elif not isinstance(base, UnboundType):
                 self.fail('Invalid base class', base_expr)
-            if isinstance(base, Instance):
-                defn.info.is_enum = self.decide_is_enum(base)
         # Add 'object' as implicit base if there is no other base class.
         if (not defn.base_types and defn.fullname != 'builtins.object'):
             obj = self.object_type()
@@ -830,11 +822,6 @@ def is_base_class(self, t: TypeInfo, s: TypeInfo) -> bool:
                     visited.add(base.type)
         return False
 
-    def decide_is_enum(self, instance: Instance) -> bool:
-        """Decide if a TypeInfo should be marked as .is_enum=True"""
-        fullname = instance.type.fullname()
-        return fullname in ENUM_BASECLASSES
-
     def analyze_metaclass(self, defn: ClassDef) -> None:
         if defn.metaclass:
             sym = self.lookup_qualified(defn.metaclass, defn)

mypy/test/data/pythoneval-enum.test

@@ -103,3 +103,18 @@ Color.m2('')
 [out]
 _program.py:11: error: Argument 1 to "m" of "Color" has incompatible type "str"; expected "int"
 _program.py:12: error: Argument 1 to "m2" of "Color" has incompatible type "str"; expected "int"
+
+[case testIntEnum_ExtendedIntEnum_functionTakingExtendedIntEnum]
+from enum import IntEnum
+class ExtendedIntEnum(IntEnum):
+    pass
+class SomeExtIntEnum(ExtendedIntEnum):
+    x = 1
+
+def takes_int(i: int):
+    pass
+takes_int(SomeExtIntEnum.x)
+
+def takes_some_ext_int_enum(s: SomeExtIntEnum):
+    pass
+takes_some_ext_int_enum(SomeExtIntEnum.x)