python · hauntsaninja · Mar 2, 2021 · Feb 9, 2021
diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py
@@ -15,6 +15,7 @@
 
 import mypy.plugin  # To avoid circular imports.
 from mypy.types import Type, Instance, LiteralType, CallableType, ProperType, get_proper_type
+from mypy.nodes import TypeInfo
 
 # Note: 'enum.EnumMeta' is deliberately excluded from this list. Classes that directly use
 # enum.EnumMeta do not necessarily automatically have the 'name' and 'value' attributes.
@@ -103,6 +104,17 @@ def _infer_value_type_with_auto_fallback(
     return ctx.default_attr_type
 
 
+def _implements_new(info: TypeInfo) -> bool:
+    """Check whether __new__ comes from enum.Enum or was implemented in a
+    subclass. In the latter case, we must infer Any as long as mypy can't infer
+    the type of _value_ from assignments in __new__.
+    """
+    type_with_new = _first(ti for ti in info.mro if ti.names.get('__new__'))
+    if type_with_new is None:
+        return False
+    return type_with_new.fullname != 'enum.Enum'
+
+
 def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:
     """This plugin refines the 'value' attribute in enums to refer to
     the original underlying value. For example, suppose we have the
@@ -135,12 +147,22 @@ class SomeEnum:
         # The value-type is still known.
         if isinstance(ctx.type, Instance):
             info = ctx.type.type
+
+            # As long as mypy doesn't understand attribute creation in __new__,
+            # there is no way to predict the value type if the enum class has a
+            # custom implementation
+            if _implements_new(info):
+                return ctx.default_attr_type
+
             stnodes = (info.get(name) for name in info.names)
-            # Enums _can_ have methods.
-            # Omit methods for our value inference.
+
+            # Enums _can_ have methods and instance attributes.
+            # Omit methods and attributes created by assigning to self.*
+            # for our value inference.
             node_types = (
                 get_proper_type(n.type) if n else None
-                for n in stnodes)
+                for n in stnodes
+                if n is None or not n.implicit)
             proper_types = (
                 _infer_value_type_with_auto_fallback(ctx, t)
                 for t in node_types
@@ -158,6 +180,13 @@ class SomeEnum:
 
     assert isinstance(ctx.type, Instance)
     info = ctx.type.type
+
+    # As long as mypy doesn't understand attribute creation in __new__,
+    # there is no way to predict the value type if the enum class has a
+    # custom implementation
+    if _implements_new(info):
+        return ctx.default_attr_type
+
     stnode = info.get(enum_field_name)
     if stnode is None:
         return ctx.default_attr_type

@@ -1243,3 +1243,59 @@ class Comparator(enum.Enum):
 
 reveal_type(Comparator.__foo__)  # N: Revealed type is 'builtins.dict[builtins.str, builtins.int]'
 [builtins fixtures/dict.pyi]
+
+[case testEnumWithInstanceAttributes]
+from enum import Enum
+class Foo(Enum):
+    def __init__(self, value: int) -> None:
+        self.foo = "bar"
+    A = 1
+    B = 2
+
+a = Foo.A
+reveal_type(a.value)    # N: Revealed type is 'builtins.int'
+reveal_type(a._value_)  # N: Revealed type is 'builtins.int'
+
+[case testNewSetsUnexpectedValueType]
+from enum import Enum
+
+class bytes:
+    def __new__(cls): pass
+
+class Foo(bytes, Enum):
+    def __new__(cls, value: int) -> 'Foo':
+        obj = bytes.__new__(cls)
+        obj._value_ = "Number %d" % value
+        return obj
+    A = 1
+    B = 2
+
+a = Foo.A
+reveal_type(a.value)    # N: Revealed type is 'Any'
+reveal_type(a._value_)  # N: Revealed type is 'Any'
+[builtins fixtures/__new__.pyi]
+[builtins fixtures/primitives.pyi]
+[typing fixtures/typing-medium.pyi]
+
+[case testValueTypeWithNewInParentClass]
+from enum import Enum
+
+class bytes:
+    def __new__(cls): pass
+
+class Foo(bytes, Enum):
+    def __new__(cls, value: int) -> 'Foo':
+        obj = bytes.__new__(cls)
+        obj._value_ = "Number %d" % value
+        return obj
+
+class Bar(Foo):
+    A = 1
+    B = 2
+
+a = Bar.A
+reveal_type(a.value)    # N: Revealed type is 'Any'
+reveal_type(a._value_)  # N: Revealed type is 'Any'
+[builtins fixtures/__new__.pyi]
+[builtins fixtures/primitives.pyi]
+[typing fixtures/typing-medium.pyi]