Add more precise inference for enum attributes

mypy/checkmember.py

@@ -615,7 +615,8 @@ def analyze_class_attribute_access(itype: Instance,
         check_final_member(name, itype.type, mx.msg, mx.context)
 
     if itype.type.is_enum and not (mx.is_lvalue or is_decorated or is_method):
-        return itype
+        enum_literal = LiteralType(name, fallback=itype)
+        return itype.copy_modified(last_known_value=enum_literal)
 
     t = node.type
     if t:

mypy/plugins/default.py

@@ -69,11 +69,16 @@ def get_method_hook(self, fullname: str
     def get_attribute_hook(self, fullname: str
                            ) -> Optional[Callable[[AttributeContext], Type]]:
         from mypy.plugins import ctypes
+        from mypy.plugins import enums
 
         if fullname == 'ctypes.Array.value':
             return ctypes.array_value_callback
         elif fullname == 'ctypes.Array.raw':
             return ctypes.array_raw_callback
+        elif fullname in enums.ENUM_NAME_ACCESS:
+            return enums.enum_name_callback
+        elif fullname in enums.ENUM_VALUE_ACCESS:
+            return enums.enum_value_callback
         return None
 
     def get_class_decorator_hook(self, fullname: str

mypy/plugins/enums.py

@@ -0,0 +1,130 @@
+"""
+This file contains a variety of plugins for refining how mypy infers types of
+expressions involving Enums.
+
+Currently, this file focuses on providing better inference for expressions like
+'SomeEnum.FOO.name' and 'SomeEnum.FOO.value'. Note that the type of both expressions
+will vary depending on exactly which instance of SomeEnum we're looking at.
+
+Note that this file does *not* contain all special-cased logic related to enums:
+we actually bake some of it directly in to the semantic analysis layer (see
+semanal_enum.py).
+"""
+from typing import Optional
+MYPY = False
+if MYPY:
+    from typing_extensions import Final
+import mypy.plugin  # To avoid circular imports.
+from mypy.types import Type, Instance, LiteralType
+
+# 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.
+ENUM_PREFIXES = {'enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag'}  # type: Final
+ENUM_NAME_ACCESS = (
+    {'{}.name'.format(prefix) for prefix in ENUM_PREFIXES}
+    | {'{}._name_'.format(prefix) for prefix in ENUM_PREFIXES}
+)  # type: Final
+ENUM_VALUE_ACCESS = (
+    {'{}.value'.format(prefix) for prefix in ENUM_PREFIXES}
+    | {'{}._value_'.format(prefix) for prefix in ENUM_PREFIXES}
+)  # type: Final
+
+
+def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type:
+    """This plugin refines the 'name' attribute in enums to act as if
+    they were declared to be final.
+
+    For example, the expression 'MyEnum.FOO.name' normally is inferred
+    to be of type 'str'.
+
+    This plugin will instead make the inferred type be a 'str' where the
+    last known value is 'Literal["FOO"]'. This means it would be legal to
+    use 'MyEnum.FOO.name' in contexts that expect a Literal type, just like
+    any other Final variable or attribute.
+
+    This plugin assumes that the provided context is an attribute access
+    matching one of the strings found in 'ENUM_NAME_ACCESS'.
+    """
+    enum_field_name = _extract_underlying_field_name(ctx.type)
+    if enum_field_name is None:
+        return ctx.default_attr_type
+    else:
+        str_type = ctx.api.named_generic_type('builtins.str', [])
+        literal_type = LiteralType(enum_field_name, fallback=str_type)
+        return str_type.copy_modified(last_known_value=literal_type)
+
+
+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
+    following:
+
+        class SomeEnum:
+            FOO = A()
+            BAR = B()
+
+    By default, mypy will infer that 'SomeEnum.FOO.value' and
+    'SomeEnum.BAR.value' both are of type 'Any'. This plugin refines
+    this inference so that mypy understands the expressions are
+    actually of types 'A' and 'B' respectively. This better reflects
+    the actual runtime behavior.
+
+    This plugin works simply by looking up the original value assigned
+    to the enum. For example, when this plugin sees 'SomeEnum.BAR.value',
+    it will look up whatever type 'BAR' had in the SomeEnum TypeInfo and
+    use that as the inferred type of the overall expression.
+
+    This plugin assumes that the provided context is an attribute access
+    matching one of the strings found in 'ENUM_VALUE_ACCESS'.
+    """
+    enum_field_name = _extract_underlying_field_name(ctx.type)
+    if enum_field_name is None:
+        return ctx.default_attr_type
+
+    assert isinstance(ctx.type, Instance)
+    info = ctx.type.type
+    stnode = info.get(enum_field_name)
+    if stnode is None:
+        return ctx.default_attr_type
+
+    underlying_type = stnode.type
+    if underlying_type is None:
+        # TODO: Deduce the inferred type if the user omits adding their own default types.
+        # TODO: Consider using the return type of `Enum._generate_next_value_` here?
+        return ctx.default_attr_type
+
+    if isinstance(underlying_type, Instance) and underlying_type.type.fullname() == 'enum.auto':
+        # TODO: Deduce the correct inferred type when the user uses 'enum.auto'.
+        # We should use the same strategy we end up picking up above.
+        return ctx.default_attr_type
+
+    return underlying_type
+
+
+def _extract_underlying_field_name(typ: Type) -> Optional[str]:
+    """If the given type corresponds to some Enum instance, returns the
+    original name of that enum. For example, if we receive in the type
+    corresponding to 'SomeEnum.FOO', we return the string "SomeEnum.Foo".
+
+    This helper takes advantage of the fact that Enum instances are valid
+    to use inside Literal[...] types. An expression like 'SomeEnum.FOO' is
+    actually represented by an Instance type with a Literal enum fallback.
+
+    We can examine this Literal fallback to retrieve the string.
+    """
+
+    if not isinstance(typ, Instance):
+        return None
+
+    if not typ.type.is_enum:
+        return None
+
+    underlying_literal = typ.last_known_value
+    if underlying_literal is None:
+        return None
+
+    # The checks above have verified this LiteralType is representing an enum value,
+    # which means the 'value' field is guaranteed to be the name of the enum field
+    # as a string.
+    assert isinstance(underlying_literal.value, str)
+    return underlying_literal.value

test-data/unit/check-enum.test

@@ -53,12 +53,9 @@ class Truth(Enum):
     false = False
 x = ''
 x = Truth.true.name
-reveal_type(Truth.true.name)
-reveal_type(Truth.false.value)
+reveal_type(Truth.true.name)    # E: Revealed type is 'builtins.str'
+reveal_type(Truth.false.value)  # E: Revealed type is 'builtins.bool'
 [builtins fixtures/bool.pyi]
-[out]
-main:7: error: Revealed type is 'builtins.str'
-main:8: error: Revealed type is 'Any'
 
 [case testEnumUnique]
 import enum
@@ -299,7 +296,7 @@ reveal_type(F.bar.name)
 [out]
 main:4: error: Revealed type is '__main__.E'
 main:5: error: Revealed type is '__main__.F'
-main:6: error: Revealed type is 'Any'
+main:6: error: Revealed type is 'builtins.int'
 main:7: error: Revealed type is 'builtins.str'
 
 [case testFunctionalEnumDict]
@@ -313,7 +310,7 @@ reveal_type(F.bar.name)
 [out]
 main:4: error: Revealed type is '__main__.E'
 main:5: error: Revealed type is '__main__.F'
-main:6: error: Revealed type is 'Any'
+main:6: error: Revealed type is 'builtins.int'
 main:7: error: Revealed type is 'builtins.str'
 
 [case testFunctionalEnumErrors]
@@ -366,11 +363,14 @@ main:22: error: "Type[W]" has no attribute "c"
 from enum import Flag, IntFlag
 A = Flag('A', 'x y')
 B = IntFlag('B', 'a b')
-reveal_type(A.x)
-reveal_type(B.a)
-[out]
-main:4: error: Revealed type is '__main__.A'
-main:5: error: Revealed type is '__main__.B'
+reveal_type(A.x)        # E: Revealed type is '__main__.A'
+reveal_type(B.a)        # E: Revealed type is '__main__.B'
+reveal_type(A.x.name)   # E: Revealed type is 'builtins.str'
+reveal_type(B.a.name)   # E: Revealed type is 'builtins.str'
+
+# TODO: The revealed type should be 'int' here
+reveal_type(A.x.value)  # E: Revealed type is 'Any'
+reveal_type(B.a.value)  # E: Revealed type is 'Any'
 
 [case testAnonymousFunctionalEnum]
 from enum import Enum
@@ -456,3 +456,157 @@ main:3: error: Revealed type is 'm.F'
 [out2]
 main:2: error: Revealed type is 'm.E'
 main:3: error: Revealed type is 'm.F'
+
+[case testEnumAuto]
+from enum import Enum, auto
+class Test(Enum):
+    a = auto()
+    b = auto()
+
+reveal_type(Test.a)  # E: Revealed type is '__main__.Test'
+[builtins fixtures/primitives.pyi]
+
+[case testEnumAttributeAccessMatrix]
+from enum import Enum, IntEnum, IntFlag, Flag, EnumMeta, auto
+from typing_extensions import Literal
+
+def is_x(val: Literal['x']) -> None: pass
+
+A1 = Enum('A1', 'x')
+class A2(Enum):
+    x = auto()
+class A3(Enum):
+    x = 1
+
+is_x(reveal_type(A1.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(A1.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(A1.x.value)         # E: Revealed type is 'Any'
+reveal_type(A1.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(A2.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(A2.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(A2.x.value)         # E: Revealed type is 'Any'
+reveal_type(A2.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(A3.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(A3.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(A3.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(A3.x._value_)       # E: Revealed type is 'builtins.int'
+
+B1 = IntEnum('B1', 'x')
+class B2(IntEnum):
+    x = auto()
+class B3(IntEnum):
+    x = 1
+
+# TODO: getting B1.x._value_ and B2.x._value_ to have type 'int' requires a typeshed change
+
+is_x(reveal_type(B1.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(B1.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(B1.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(B1.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(B2.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(B2.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(B2.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(B2.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(B3.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(B3.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(B3.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(B3.x._value_)       # E: Revealed type is 'builtins.int'
+
+# TODO: C1.x.value and C2.x.value should also be of type 'int'
+# This requires either a typeshed change or a plugin refinement
+
+C1 = IntFlag('C1', 'x')
+class C2(IntFlag):
+    x = auto()
+class C3(IntFlag):
+    x = 1
+
+is_x(reveal_type(C1.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(C1.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(C1.x.value)         # E: Revealed type is 'Any'
+reveal_type(C1.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(C2.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(C2.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(C2.x.value)         # E: Revealed type is 'Any'
+reveal_type(C2.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(C3.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(C3.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(C3.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(C3.x._value_)       # E: Revealed type is 'builtins.int'
+
+D1 = Flag('D1', 'x')
+class D2(Flag):
+    x = auto()
+class D3(Flag):
+    x = 1
+
+is_x(reveal_type(D1.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(D1.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(D1.x.value)         # E: Revealed type is 'Any'
+reveal_type(D1.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(D2.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(D2.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(D2.x.value)         # E: Revealed type is 'Any'
+reveal_type(D2.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(D3.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(D3.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(D3.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(D3.x._value_)       # E: Revealed type is 'builtins.int'
+
+# TODO: Generalize our enum functional API logic to work with subclasses of Enum
+# See https://github.com/python/mypy/issues/6037
+
+class Parent(Enum): pass
+# E1 = Parent('E1', 'x')  # See above TODO
+class E2(Parent):
+    x = auto()
+class E3(Parent):
+    x = 1
+
+is_x(reveal_type(E2.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(E2.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(E2.x.value)         # E: Revealed type is 'Any'
+reveal_type(E2.x._value_)       # E: Revealed type is 'Any'
+is_x(reveal_type(E3.x.name))    # E: Revealed type is 'Literal['x']'
+is_x(reveal_type(E3.x._name_))  # E: Revealed type is 'Literal['x']'
+reveal_type(E3.x.value)         # E: Revealed type is 'builtins.int'
+reveal_type(E3.x._value_)       # E: Revealed type is 'builtins.int'
+
+
+# TODO: Figure out if we can construct enums using EnumMetas using the functional API.
+# Also figure out if we even care about supporting that use case.
+class F2(metaclass=EnumMeta):
+    x = auto()
+class F3(metaclass=EnumMeta):
+    x = 1
+
+F2.x.name      # E: "F2" has no attribute "name"
+F2.x._name_    # E: "F2" has no attribute "_name_"
+F2.x.value     # E: "F2" has no attribute "value"
+F2.x._value_   # E: "F2" has no attribute "_value_"
+F3.x.name      # E: "F3" has no attribute "name"
+F3.x._name_    # E: "F3" has no attribute "_name_"
+F3.x.value     # E: "F3" has no attribute "value"
+F3.x._value_   # E: "F3" has no attribute "_value_"
+[builtins fixtures/primitives.pyi]
+
+[case testEnumAttributeChangeIncremental]
+from a import SomeEnum
+reveal_type(SomeEnum.a.value)
+
+[file a.py]
+from b import SomeEnum
+
+[file b.py]
+from enum import Enum
+class SomeEnum(Enum):
+    a = 1
+
+[file b.py.2]
+from enum import Enum
+class SomeEnum(Enum):
+    a = "foo"
+[out]
+main:2: error: Revealed type is 'builtins.int'
+[out2]
+main:2: error: Revealed type is 'builtins.str'