From 815a7e35a5170300c36ee19368be7bc8238ba292 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Sun, 12 May 2019 07:51:02 -0700 Subject: [PATCH 1/6] Add more precise inference for enum attributes This pull request makes two changes to enum attributes. First, this PR refines type inference for expressions like `MyEnum.FOO` and `MyEnum.FOO.name`. Those two expressions will continue to evaluate to `MyEnum` and `str` respectively under normal conditions, but will evaluate to `Literal[MyEnum.FOO]` and `Literal["FOO"]` respectively when used in Literal contexts. Second, the type of `MyEnum.FOO.value` will be more precise when possible: mypy will evaluate that expression to the type of whatever FOO was assigned in the enum definition, falling back to `Any` as a default. Somewhat relatedly, this diff adds a few tests confirming we handle enum.auto() correctly. Two additional notes: 1. The changes I made to the `name` and `value` fields up above are strictly speaking unsafe. While those files are normally read-only (doing `MyEnum.FOO.name = blah` is a runtime error), it's actually possible to change those fields anyway by altering the `_name_` and `_value_` fields which are *not* protected. But I think this use case is probably rare -- I'm planning on investigating the feasibility of just having mypy just disallow modifying these attributes altogether after I investigate how enums are used in some internal codebases in a little more detail. 2. I would have liked to make `MyEnum.FOO.value` also return an even more precise type when used in literal contexts similar to `MyEnum.FOO.name`, but I think our plugin system needs to be a bit more flexible first. --- mypy/checkmember.py | 3 +- mypy/plugins/default.py | 5 + mypy/plugins/enums.py | 69 +++++++++++ test-data/unit/check-enum.test | 175 +++++++++++++++++++++++++-- test-data/unit/check-literal.test | 126 ++++++++++++++----- test-data/unit/lib-stub/builtins.pyi | 1 + test-data/unit/lib-stub/enum.pyi | 14 ++- 7 files changed, 346 insertions(+), 47 deletions(-) create mode 100644 mypy/plugins/enums.py diff --git a/mypy/checkmember.py b/mypy/checkmember.py index e17d5f044d844..e14d3ce97b72d 100644 --- a/mypy/checkmember.py +++ b/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: diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 059feab1d6a25..fb0177267257a 100644 --- a/mypy/plugins/default.py +++ b/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 diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py new file mode 100644 index 0000000000000..37051d6a5e423 --- /dev/null +++ b/mypy/plugins/enums.py @@ -0,0 +1,69 @@ +from typing import Optional +import mypy.plugin # To avoid circular imports. +from mypy.types import Type, Instance, LiteralType +from mypy.nodes import Var, MDEF + +# 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'] +ENUM_NAME_ACCESS = ( + ['{}.name'.format(prefix) for prefix in ENUM_PREFIXES] + + ['{}._name_'.format(prefix) for prefix in ENUM_PREFIXES] +) +ENUM_VALUE_ACCESS = ( + ['{}.value'.format(prefix) for prefix in ENUM_PREFIXES] + + ['{}._value_'.format(prefix) for prefix in ENUM_PREFIXES] +) + + +def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: + 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: + 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 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 diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 8a89d35ba9409..e83fa3e6ff15d 100644 --- a/test-data/unit/check-enum.test +++ b/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,154 @@ 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' + +[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 + +class Parent(Enum): pass +#E1 = Parent('E1', 'x') +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_" + +[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' diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test index 3ad0b7486d2a6..a70fb2e33626b 100644 --- a/test-data/unit/check-literal.test +++ b/test-data/unit/check-literal.test @@ -2651,6 +2651,45 @@ d_wrap: Literal[4, d] # E: Invalid type "__main__.d" \ # E: Parameter 2 of Literal[...] is invalid [out] +[case testLiteralWithFinalPropagation] +from typing_extensions import Final, Literal + +a: Final = 3 +b: Final = a +c = a + +def expect_3(x: Literal[3]) -> None: pass +expect_3(a) +expect_3(b) +expect_3(c) # E: Argument 1 to "expect_3" has incompatible type "int"; expected "Literal[3]" +[out] + +[case testLiteralWithFinalPropagationIsNotLeaking] +from typing_extensions import Final, Literal + +final_tuple_direct: Final = (2, 3) +final_tuple_indirect: Final = final_tuple_direct +mutable_tuple = final_tuple_direct +final_list_1: Final = [2] +final_list_2: Final = [2, 2] +final_dict: Final = {"foo": 2} +final_set_1: Final = {2} +final_set_2: Final = {2, 2} + +def expect_2(x: Literal[2]) -> None: pass + +expect_2(final_tuple_direct[0]) +expect_2(final_tuple_indirect[0]) + +expect_2(mutable_tuple[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +expect_2(final_list_1[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +expect_2(final_list_2[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +expect_2(final_dict["foo"]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +expect_2(final_set_1.pop()) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +expect_2(final_set_2.pop()) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" +[builtins fixtures/isinstancelist.pyi] +[out] + -- -- Tests for Literals and enums -- @@ -2814,41 +2853,68 @@ x: Literal[Alias.FOO] reveal_type(x) # E: Revealed type is 'Literal[__main__.Test.FOO]' [out] -[case testLiteralWithFinalPropagation] -from typing_extensions import Final, Literal +[case testLiteralUsingEnumAttributesInLiteralContexts] +from typing_extensions import Literal, Final +from enum import Enum -a: Final = 3 -b: Final = a -c = a +class Test1(Enum): + FOO = 1 + BAR = 2 +Test2 = Enum('Test2', [('FOO', 1), ('BAR', 2)]) -def expect_3(x: Literal[3]) -> None: pass -expect_3(a) -expect_3(b) -expect_3(c) # E: Argument 1 to "expect_3" has incompatible type "int"; expected "Literal[3]" -[out] +def expects_test1_foo(x: Literal[Test1.FOO]) -> None: ... +def expects_test2_foo(x: Literal[Test2.FOO]) -> None: ... -[case testLiteralWithFinalPropagationIsNotLeaking] -from typing_extensions import Final, Literal +expects_test1_foo(Test1.FOO) +expects_test1_foo(Test1.BAR) # E: Argument 1 to "expects_test1_foo" has incompatible type "Literal[Test1.BAR]"; expected "Literal[Test1.FOO]" +expects_test2_foo(Test2.FOO) +expects_test2_foo(Test2.BAR) # E: Argument 1 to "expects_test2_foo" has incompatible type "Literal[Test2.BAR]"; expected "Literal[Test2.FOO]" -final_tuple_direct: Final = (2, 3) -final_tuple_indirect: Final = final_tuple_direct -mutable_tuple = final_tuple_direct -final_list_1: Final = [2] -final_list_2: Final = [2, 2] -final_dict: Final = {"foo": 2} -final_set_1: Final = {2} -final_set_2: Final = {2, 2} +# Make sure the two 'FOO's are not interchangeable +expects_test1_foo(Test2.FOO) # E: Argument 1 to "expects_test1_foo" has incompatible type "Literal[Test2.FOO]"; expected "Literal[Test1.FOO]" +expects_test2_foo(Test1.FOO) # E: Argument 1 to "expects_test2_foo" has incompatible type "Literal[Test1.FOO]"; expected "Literal[Test2.FOO]" -def expect_2(x: Literal[2]) -> None: pass +# Make sure enums follow the same semantics as 'x = 1' vs 'x: Final = 1' +var1 = Test1.FOO +final1: Final = Test1.FOO +expects_test1_foo(var1) # E: Argument 1 to "expects_test1_foo" has incompatible type "Test1"; expected "Literal[Test1.FOO]" +expects_test1_foo(final1) -expect_2(final_tuple_direct[0]) -expect_2(final_tuple_indirect[0]) +var2 = Test2.FOO +final2: Final = Test2.FOO +expects_test2_foo(var2) # E: Argument 1 to "expects_test2_foo" has incompatible type "Test2"; expected "Literal[Test2.FOO]" +expects_test2_foo(final2) +[out] -expect_2(mutable_tuple[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -expect_2(final_list_1[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -expect_2(final_list_2[0]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -expect_2(final_dict["foo"]) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -expect_2(final_set_1.pop()) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -expect_2(final_set_2.pop()) # E: Argument 1 to "expect_2" has incompatible type "int"; expected "Literal[2]" -[builtins fixtures/isinstancelist.pyi] +[case testLiteralUsingEnumAttributeNamesInLiteralContexts] +from typing_extensions import Literal, Final +from enum import Enum + +class Test1(Enum): + FOO = 1 + BAR = 2 +Test2 = Enum('Test2', [('FOO', 1), ('BAR', 2)]) +Test3 = Enum('Test3', 'FOO BAR') +Test4 = Enum('Test4', ['FOO', 'BAR']) +Test5 = Enum('Test5', {'FOO': 1, 'BAR': 2}) + +def expects_foo(x: Literal['FOO']) -> None: ... + +expects_foo(Test1.FOO.name) +expects_foo(Test2.FOO.name) +expects_foo(Test3.FOO.name) +expects_foo(Test4.FOO.name) +expects_foo(Test5.FOO.name) + +expects_foo(Test1.BAR.name) # E: Argument 1 to "expects_foo" has incompatible type "Literal['BAR']"; expected "Literal['FOO']" +expects_foo(Test2.BAR.name) # E: Argument 1 to "expects_foo" has incompatible type "Literal['BAR']"; expected "Literal['FOO']" +expects_foo(Test3.BAR.name) # E: Argument 1 to "expects_foo" has incompatible type "Literal['BAR']"; expected "Literal['FOO']" +expects_foo(Test4.BAR.name) # E: Argument 1 to "expects_foo" has incompatible type "Literal['BAR']"; expected "Literal['FOO']" +expects_foo(Test5.BAR.name) # E: Argument 1 to "expects_foo" has incompatible type "Literal['BAR']"; expected "Literal['FOO']" + +reveal_type(Test1.FOO.name) # E: Revealed type is 'builtins.str' +reveal_type(Test2.FOO.name) # E: Revealed type is 'builtins.str' +reveal_type(Test3.FOO.name) # E: Revealed type is 'builtins.str' +reveal_type(Test4.FOO.name) # E: Revealed type is 'builtins.str' +reveal_type(Test5.FOO.name) # E: Revealed type is 'builtins.str' [out] diff --git a/test-data/unit/lib-stub/builtins.pyi b/test-data/unit/lib-stub/builtins.pyi index 29365d0efd196..6abc1b5f5537a 100644 --- a/test-data/unit/lib-stub/builtins.pyi +++ b/test-data/unit/lib-stub/builtins.pyi @@ -9,6 +9,7 @@ class type: # These are provided here for convenience. class int: + def __init__(self, x: object = ..., base: int = ...) -> None: pass def __add__(self, other: int) -> int: pass def __rmul__(self, other: int) -> int: pass class float: pass diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi index 7b97b5640a0e0..14908c2d10635 100644 --- a/test-data/unit/lib-stub/enum.pyi +++ b/test-data/unit/lib-stub/enum.pyi @@ -9,18 +9,20 @@ class EnumMeta(type, Sized): def __getitem__(self: Type[_T], name: str) -> _T: pass class Enum(metaclass=EnumMeta): - def __new__(cls, value: Any) -> None: pass + def __new__(cls: Type[_T], value: object) -> _T: pass def __repr__(self) -> str: pass def __str__(self) -> str: pass def __format__(self, format_spec: str) -> str: pass def __hash__(self) -> Any: pass def __reduce_ex__(self, proto: Any) -> Any: pass - name = '' # type: str - value = None # type: Any + name: str + value: Any + _name_: str + _value_: Any class IntEnum(int, Enum): - value = 0 # type: int + value: int def unique(enumeration: _T) -> _T: pass @@ -32,3 +34,7 @@ class Flag(Enum): class IntFlag(int, Flag): def __and__(self: _T, other: Union[int, _T]) -> _T: pass + + +class auto(IntFlag): + value: Any \ No newline at end of file From 952d2d3e4e5a2cb4fabee7607fc0b6452bdb9dff Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Tue, 21 May 2019 09:48:40 -0700 Subject: [PATCH 2/6] Fix broken test --- test-data/unit/check-protocols.test | 2 +- test-data/unit/lib-stub/builtins.pyi | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/test-data/unit/check-protocols.test b/test-data/unit/check-protocols.test index 55ca338737485..c961795f63381 100644 --- a/test-data/unit/check-protocols.test +++ b/test-data/unit/check-protocols.test @@ -2427,7 +2427,7 @@ from typing import Protocol class P(Protocol): ... class C(P): ... -reveal_type(C.register(int)) # E: Revealed type is 'def () -> builtins.int' +reveal_type(C.register(int)) # E: Revealed type is 'def (x: builtins.object =, base: builtins.int =) -> builtins.int' [typing fixtures/typing-full.pyi] [out] diff --git a/test-data/unit/lib-stub/builtins.pyi b/test-data/unit/lib-stub/builtins.pyi index 6abc1b5f5537a..d3555c33ca233 100644 --- a/test-data/unit/lib-stub/builtins.pyi +++ b/test-data/unit/lib-stub/builtins.pyi @@ -9,6 +9,7 @@ class type: # These are provided here for convenience. class int: + # Note: this is a simplification of the actual signature def __init__(self, x: object = ..., base: int = ...) -> None: pass def __add__(self, other: int) -> int: pass def __rmul__(self, other: int) -> int: pass From e14b629b9f8aa3fd5663b52a771596f9ccdc3986 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Sun, 26 May 2019 16:22:56 -0700 Subject: [PATCH 3/6] Respond to code review --- mypy/plugins/enums.py | 79 ++++++++++++++++++++++---- test-data/unit/check-enum.test | 5 +- test-data/unit/check-protocols.test | 2 +- test-data/unit/fixtures/primitives.pyi | 2 + test-data/unit/lib-stub/builtins.pyi | 2 - 5 files changed, 75 insertions(+), 15 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 37051d6a5e423..e2c30e3a130c0 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -1,23 +1,49 @@ +""" +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 +from typing_extensions import Final import mypy.plugin # To avoid circular imports. from mypy.types import Type, Instance, LiteralType -from mypy.nodes import Var, MDEF # 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'] -ENUM_NAME_ACCESS = ( - ['{}.name'.format(prefix) for prefix in ENUM_PREFIXES] - + ['{}._name_'.format(prefix) for prefix in ENUM_PREFIXES] +ENUM_PREFIXES: Final = {'enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag'} +ENUM_NAME_ACCESS: Final = ( + {'{}.name'.format(prefix) for prefix in ENUM_PREFIXES} + | {'{}._name_'.format(prefix) for prefix in ENUM_PREFIXES} ) -ENUM_VALUE_ACCESS = ( - ['{}.value'.format(prefix) for prefix in ENUM_PREFIXES] - + ['{}._value_'.format(prefix) for prefix in ENUM_PREFIXES] +ENUM_VALUE_ACCESS: Final = ( + {'{}.value'.format(prefix) for prefix in ENUM_PREFIXES} + | {'{}._value_'.format(prefix) for prefix in ENUM_PREFIXES} ) def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: - enum_field_name = extract_underlying_field_name(ctx.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: @@ -27,7 +53,27 @@ def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: - enum_field_name = extract_underlying_field_name(ctx.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, + + 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 @@ -51,7 +97,18 @@ def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: return underlying_type -def extract_underlying_field_name(typ: Type) -> Optional[str]: +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 diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index e83fa3e6ff15d..26c7609d07df4 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -464,6 +464,7 @@ class Test(Enum): 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 @@ -553,9 +554,10 @@ 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') +# E1 = Parent('E1', 'x') # See above TODO class E2(Parent): x = auto() class E3(Parent): @@ -586,6 +588,7 @@ 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 diff --git a/test-data/unit/check-protocols.test b/test-data/unit/check-protocols.test index c961795f63381..55ca338737485 100644 --- a/test-data/unit/check-protocols.test +++ b/test-data/unit/check-protocols.test @@ -2427,7 +2427,7 @@ from typing import Protocol class P(Protocol): ... class C(P): ... -reveal_type(C.register(int)) # E: Revealed type is 'def (x: builtins.object =, base: builtins.int =) -> builtins.int' +reveal_type(C.register(int)) # E: Revealed type is 'def () -> builtins.int' [typing fixtures/typing-full.pyi] [out] diff --git a/test-data/unit/fixtures/primitives.pyi b/test-data/unit/fixtures/primitives.pyi index f2c0cd03acfce..07f1daf14193d 100644 --- a/test-data/unit/fixtures/primitives.pyi +++ b/test-data/unit/fixtures/primitives.pyi @@ -12,6 +12,8 @@ class type: def __init__(self, x) -> None: pass class int: + # Note: this is a simplification of the actual signature + def __init__(self, x: object = ..., base: int = ...) -> None: pass def __add__(self, i: int) -> int: pass class float: def __float__(self) -> float: pass diff --git a/test-data/unit/lib-stub/builtins.pyi b/test-data/unit/lib-stub/builtins.pyi index d3555c33ca233..29365d0efd196 100644 --- a/test-data/unit/lib-stub/builtins.pyi +++ b/test-data/unit/lib-stub/builtins.pyi @@ -9,8 +9,6 @@ class type: # These are provided here for convenience. class int: - # Note: this is a simplification of the actual signature - def __init__(self, x: object = ..., base: int = ...) -> None: pass def __add__(self, other: int) -> int: pass def __rmul__(self, other: int) -> int: pass class float: pass From fee282707993ea8429a263ece7f300be67b20650 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Tue, 28 May 2019 08:07:48 -0700 Subject: [PATCH 4/6] Respond to code review, v2 --- mypy/plugins/enums.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index e2c30e3a130c0..1eff4a56e0d99 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -17,15 +17,15 @@ # 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: Final = {'enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag'} -ENUM_NAME_ACCESS: Final = ( +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} -) -ENUM_VALUE_ACCESS: Final = ( +) # 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: @@ -68,7 +68,9 @@ class SomeEnum: the actual runtime behavior. This plugin works simply by looking up the original value assigned - to the enum. For example, + 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'. From b85ace43d72fe74b2395e56c63e01b55be8419b9 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Tue, 28 May 2019 23:44:32 -0700 Subject: [PATCH 5/6] Experiment with adding typing_extensions as a dependency --- setup.py | 1 + test-requirements.txt | 1 + 2 files changed, 2 insertions(+) diff --git a/setup.py b/setup.py index 6bee91e5616c8..acdccd795b328 100644 --- a/setup.py +++ b/setup.py @@ -180,6 +180,7 @@ def run(self): # When changing this, also update test-requirements.txt. install_requires=['typed-ast >= 1.3.5, < 1.4.0', 'mypy_extensions >= 0.4.0, < 0.5.0', + 'typing_extensions >= 3.7.0, < 4.0.0', ], # Same here. extras_require={'dmypy': 'psutil >= 4.0'}, diff --git a/test-requirements.txt b/test-requirements.txt index 5f310d5930240..d3f7d9cbc6d78 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -4,6 +4,7 @@ flake8-bugbear; python_version >= '3.5' flake8-pyi; python_version >= '3.6' lxml==4.2.4 mypy_extensions>=0.4.0,<0.5.0 +typing_extensions>=3.7.0,<4.0.0 psutil>=4.0 pytest>=4.4 pytest-xdist>=1.22 From a8d2139c5c0af90336fa1b5b5e6715fbe5669e08 Mon Sep 17 00:00:00 2001 From: Michael Lee Date: Wed, 29 May 2019 00:24:32 -0700 Subject: [PATCH 6/6] Apply Ethan's suggestion --- mypy/plugins/enums.py | 4 +++- setup.py | 1 - test-requirements.txt | 1 - 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 1eff4a56e0d99..645afa441844a 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -11,7 +11,9 @@ semanal_enum.py). """ from typing import Optional -from typing_extensions import Final +MYPY = False +if MYPY: + from typing_extensions import Final import mypy.plugin # To avoid circular imports. from mypy.types import Type, Instance, LiteralType diff --git a/setup.py b/setup.py index acdccd795b328..6bee91e5616c8 100644 --- a/setup.py +++ b/setup.py @@ -180,7 +180,6 @@ def run(self): # When changing this, also update test-requirements.txt. install_requires=['typed-ast >= 1.3.5, < 1.4.0', 'mypy_extensions >= 0.4.0, < 0.5.0', - 'typing_extensions >= 3.7.0, < 4.0.0', ], # Same here. extras_require={'dmypy': 'psutil >= 4.0'}, diff --git a/test-requirements.txt b/test-requirements.txt index d3f7d9cbc6d78..5f310d5930240 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -4,7 +4,6 @@ flake8-bugbear; python_version >= '3.5' flake8-pyi; python_version >= '3.6' lxml==4.2.4 mypy_extensions>=0.4.0,<0.5.0 -typing_extensions>=3.7.0,<4.0.0 psutil>=4.0 pytest>=4.4 pytest-xdist>=1.22