From fa67a318f1ec3d3df9c2cbd422d7af1dd1a34c26 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Tue, 15 Sep 2020 09:57:23 -0400 Subject: [PATCH 01/13] Predict enum value type for unknown member names. It is very common for enums to have homogenous member-value types. In the case where we do not know what enum member we are dealing with, we should sniff for that case and still collapse to a known type if that assuption holds. --- mypy/plugins/enums.py | 16 ++++++++++++++++ test-data/unit/check-enum.test | 20 ++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 81aa29afcb11..959850accad5 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -78,6 +78,22 @@ class SomeEnum: """ enum_field_name = _extract_underlying_field_name(ctx.type) if enum_field_name is None: + # We do not know the ennum field name (perhaps it was passed to a function and we only + # know that it _is_ a member). All is not lost however, if we can prove that the all + # of the enum members have the same value-type, then it doesn't matter which member + # was passed in. The value-type is still known. + if isinstance(ctx.type, Instance): + info = ctx.type.type + stnodes = (info.get(name) for name in info.names) + first_node = next(stnodes, None) + if first_node is None: + return ctx.default_attr_type + first_node_type = first_node.type + if all(node is not None and node.type == first_node_type for node in stnodes): + underlying_type = get_proper_type(first_node_type) + if underlying_type is not None: + return underlying_type + return ctx.default_attr_type assert isinstance(ctx.type, Instance) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index e66fdfe277a1..dfe6f1dde646 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -59,6 +59,26 @@ reveal_type(Truth.true.name) # N: Revealed type is 'Literal['true']?' reveal_type(Truth.false.value) # N: Revealed type is 'builtins.bool' [builtins fixtures/bool.pyi] +[case testEnumValueExtended] +from enum import Enum +class Truth(Enum): + true = True + false = False + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'builtins.bool' +[builtins fixtures/bool.pyi] + +[case testEnumValueInhomogenous] +from enum import Enum +class Truth(Enum): + true = 'True' + false = 0 + +def cannot_infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'Any' +[builtins fixtures/bool.pyi] + [case testEnumUnique] import enum @enum.unique From 077d3e77cf0e3a0eff4d739f45876fa070b8add0 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Tue, 15 Sep 2020 10:28:26 -0400 Subject: [PATCH 02/13] Fix a typo. Squash before merging. --- mypy/plugins/enums.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 959850accad5..7dbefa5c1554 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -78,7 +78,7 @@ class SomeEnum: """ enum_field_name = _extract_underlying_field_name(ctx.type) if enum_field_name is None: - # We do not know the ennum field name (perhaps it was passed to a function and we only + # We do not know the enum field name (perhaps it was passed to a function and we only # know that it _is_ a member). All is not lost however, if we can prove that the all # of the enum members have the same value-type, then it doesn't matter which member # was passed in. The value-type is still known. From e86111748694a2f245278ce76de3583f48fa1a14 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 01:14:58 -0400 Subject: [PATCH 03/13] Get enum-value inference to work with auto. --- mypy/plugins/enums.py | 80 ++++++++++++++++++++++++++-------- test-data/unit/check-enum.test | 52 ++++++++++++++++++++++ 2 files changed, 115 insertions(+), 17 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 7dbefa5c1554..b5dfa459d075 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -14,7 +14,7 @@ from typing_extensions import Final import mypy.plugin # To avoid circular imports. -from mypy.types import Type, Instance, LiteralType, get_proper_type +from mypy.types import Type, Instance, LiteralType, CallableType, ProperType, get_proper_type # 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. @@ -53,6 +53,46 @@ def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: return str_type.copy_modified(last_known_value=literal_type) +def _infer_value_type_with_auto_fallback( + ctx: 'mypy.plugin.AttributeContext', + proper_type: Optional[ProperType]) -> Optional[Type]: + """Figure out the type of an enum value accounting for `auto()`. + + This method is a no-op for a `None` proper_type and also in the case where + the type is not "enum.auto" + """ + if proper_type is None: + return None + if (isinstance(proper_type, Instance) and + proper_type.type.fullname == 'enum.auto'): + info = ctx.type.type + # Find the first _generate_next_value_ on the mro. We need to know + # if it is `Enum` because `Enum` types say that the return-value of + #`_generate_next_value_` is `Any`. In reality the default `auto()` + # returns an `int` (presumably the `Any` in typeshed is to make it + # easier to subclass and change the returned type). + type_with_generate_next_value = next( + (type_info for type_info in info.mro + if type_info.names.get('_generate_next_value_')), + None) + if type_with_generate_next_value is None: + return ctx.default_attr_type + + stnode = type_with_generate_next_value.get('_generate_next_value_') + if stnode is None: + return ctx.default_attr_type + + # This should be a `CallableType` + node_type = stnode.type + if isinstance(node_type, CallableType): + if type_with_generate_next_value.fullname == 'enum.Enum': + int_type = ctx.api.named_generic_type('builtins.int', []) + return int_type + return get_proper_type(node_type.ret_type) + return ctx.default_attr_type + return proper_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 @@ -78,22 +118,33 @@ class SomeEnum: """ enum_field_name = _extract_underlying_field_name(ctx.type) if enum_field_name is None: - # We do not know the enum field name (perhaps it was passed to a function and we only - # know that it _is_ a member). All is not lost however, if we can prove that the all - # of the enum members have the same value-type, then it doesn't matter which member - # was passed in. The value-type is still known. + # We do not know the enum field name (perhaps it was passed to a + # function and we only know that it _is_ a member). All is not lost + # however, if we can prove that the all of the enum members have the + # same value-type, then it doesn't matter which member was passed in. + # The value-type is still known. if isinstance(ctx.type, Instance): info = ctx.type.type stnodes = (info.get(name) for name in info.names) - first_node = next(stnodes, None) - if first_node is None: + # Enums _can_ have methods. + # Omit methods for our value inference. + stnodes_non_method = ( + n for n in stnodes if not isinstance(n.type, CallableType)) + node_types = ( + get_proper_type(n.type) if n else None + for n in stnodes_non_method) + proper_types = ( + _infer_value_type_with_auto_fallback(ctx, t) + for t in node_types) + underlying_type = next(proper_types, None) + if underlying_type is None: return ctx.default_attr_type - first_node_type = first_node.type - if all(node is not None and node.type == first_node_type for node in stnodes): - underlying_type = get_proper_type(first_node_type) + all_same_value_type = all( + proper_type is not None and proper_type == underlying_type + for proper_type in proper_types) + if all_same_value_type: if underlying_type is not None: return underlying_type - return ctx.default_attr_type assert isinstance(ctx.type, Instance) @@ -108,12 +159,7 @@ class SomeEnum: # 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 + return _infer_value_type_with_auto_fallback(ctx, underlying_type) def _extract_underlying_field_name(typ: Type) -> Optional[str]: diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index dfe6f1dde646..328737feca63 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -69,6 +69,58 @@ def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.bool' [builtins fixtures/bool.pyi] +[case testEnumValueAllAuto] +from enum import Enum, auto +class Truth(Enum): + true = auto() + false = auto() + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'builtins.int' +[builtins fixtures/bool.pyi] +[builtins fixtures/primitives.pyi] + +[case testEnumValueSomeAuto] +from enum import Enum, auto +class Truth(Enum): + true = 8675309 + false = auto() + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'builtins.int' +[builtins fixtures/bool.pyi] + +[case testEnumValueExtraMethods] +from enum import Enum, auto +class Truth(Enum): + true = True + false = False + + def foo(self) -> str: + return 'bar' + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'builtins.bool' +[builtins fixtures/bool.pyi] + +[case testEnumValueCustomAuto] +from enum import Enum, auto +from typing import List, Any +class AutoName(Enum): + @staticmethod + def _generate_next_value_(name: str, start: int, count: int, last_values: List[Any]) -> str: + return name + +class Truth(AutoName): + true = auto() + false = auto() + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is 'builtins.str' +[builtins fixtures/bool.pyi] +[builtins fixtures/staticmethod.pyi] +[builtins fixtures/list.pyi] + [case testEnumValueInhomogenous] from enum import Enum class Truth(Enum): From db168713d0a53479de0e1ba6bd0931aed06bfa4c Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 07:46:16 -0400 Subject: [PATCH 04/13] Fix flake8 complaint. --- mypy/plugins/enums.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index b5dfa459d075..ef513301739b 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -68,7 +68,7 @@ def _infer_value_type_with_auto_fallback( info = ctx.type.type # Find the first _generate_next_value_ on the mro. We need to know # if it is `Enum` because `Enum` types say that the return-value of - #`_generate_next_value_` is `Any`. In reality the default `auto()` + # `_generate_next_value_` is `Any`. In reality the default `auto()` # returns an `int` (presumably the `Any` in typeshed is to make it # easier to subclass and change the returned type). type_with_generate_next_value = next( From bfee76d3e8a7970241124ec34e6294ea84154268 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 09:06:12 -0400 Subject: [PATCH 05/13] Fixup the unit-tests. --- test-data/unit/check-enum.test | 28 ++++++++++++++-------------- test-data/unit/lib-stub/enum.pyi | 9 ++++++++- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 328737feca63..e4e9e4bee0bd 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -77,7 +77,6 @@ class Truth(Enum): def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.int' -[builtins fixtures/bool.pyi] [builtins fixtures/primitives.pyi] [case testEnumValueSomeAuto] @@ -88,7 +87,6 @@ class Truth(Enum): def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.int' -[builtins fixtures/bool.pyi] [case testEnumValueExtraMethods] from enum import Enum, auto @@ -107,9 +105,11 @@ def infer_truth(truth: Truth) -> None: from enum import Enum, auto from typing import List, Any class AutoName(Enum): - @staticmethod - def _generate_next_value_(name: str, start: int, count: int, last_values: List[Any]) -> str: - return name + + # In `typeshed`, this is a staticmethod and has more arguments, + # but I have lied a bit to keep the test stubs lean. + def _generate_next_value_(self) -> str: + return "name" class Truth(AutoName): true = auto() @@ -569,8 +569,8 @@ reveal_type(A1.x.value) # N: Revealed type is 'Any' reveal_type(A1.x._value_) # N: Revealed type is 'Any' is_x(reveal_type(A2.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(A2.x._name_)) # N: Revealed type is 'Literal['x']' -reveal_type(A2.x.value) # N: Revealed type is 'Any' -reveal_type(A2.x._value_) # N: Revealed type is 'Any' +reveal_type(A2.x.value) # N: Revealed type is 'builtins.int' +reveal_type(A2.x._value_) # N: Revealed type is 'builtins.int' is_x(reveal_type(A3.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(A3.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(A3.x.value) # N: Revealed type is 'builtins.int' @@ -591,7 +591,7 @@ reveal_type(B1.x._value_) # N: Revealed type is 'Any' is_x(reveal_type(B2.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(B2.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(B2.x.value) # N: Revealed type is 'builtins.int' -reveal_type(B2.x._value_) # N: Revealed type is 'Any' +reveal_type(B2.x._value_) # N: Revealed type is 'builtins.int' is_x(reveal_type(B3.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(B3.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(B3.x.value) # N: Revealed type is 'builtins.int' @@ -612,8 +612,8 @@ reveal_type(C1.x.value) # N: Revealed type is 'Any' reveal_type(C1.x._value_) # N: Revealed type is 'Any' is_x(reveal_type(C2.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(C2.x._name_)) # N: Revealed type is 'Literal['x']' -reveal_type(C2.x.value) # N: Revealed type is 'Any' -reveal_type(C2.x._value_) # N: Revealed type is 'Any' +reveal_type(C2.x.value) # N: Revealed type is 'builtins.int' +reveal_type(C2.x._value_) # N: Revealed type is 'builtins.int' is_x(reveal_type(C3.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(C3.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(C3.x.value) # N: Revealed type is 'builtins.int' @@ -631,8 +631,8 @@ reveal_type(D1.x.value) # N: Revealed type is 'Any' reveal_type(D1.x._value_) # N: Revealed type is 'Any' is_x(reveal_type(D2.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(D2.x._name_)) # N: Revealed type is 'Literal['x']' -reveal_type(D2.x.value) # N: Revealed type is 'Any' -reveal_type(D2.x._value_) # N: Revealed type is 'Any' +reveal_type(D2.x.value) # N: Revealed type is 'builtins.int' +reveal_type(D2.x._value_) # N: Revealed type is 'builtins.int' is_x(reveal_type(D3.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(D3.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(D3.x.value) # N: Revealed type is 'builtins.int' @@ -650,8 +650,8 @@ class E3(Parent): is_x(reveal_type(E2.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(E2.x._name_)) # N: Revealed type is 'Literal['x']' -reveal_type(E2.x.value) # N: Revealed type is 'Any' -reveal_type(E2.x._value_) # N: Revealed type is 'Any' +reveal_type(E2.x.value) # N: Revealed type is 'builtins.int' +reveal_type(E2.x._value_) # N: Revealed type is 'builtins.int' is_x(reveal_type(E3.x.name)) # N: Revealed type is 'Literal['x']' is_x(reveal_type(E3.x._name_)) # N: Revealed type is 'Literal['x']' reveal_type(E3.x.value) # N: Revealed type is 'builtins.int' diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi index 14908c2d1063..29178732ea32 100644 --- a/test-data/unit/lib-stub/enum.pyi +++ b/test-data/unit/lib-stub/enum.pyi @@ -21,6 +21,10 @@ class Enum(metaclass=EnumMeta): _name_: str _value_: Any + # In reality, _generate_next_value_ is python3.6 only and has a different signature. + # However, this should be quick and doesn't require additional stubs (e.g. `staticmethod`) + def _generate_next_value_(self) -> Any: pass + class IntEnum(int, Enum): value: int @@ -37,4 +41,7 @@ class IntFlag(int, Flag): class auto(IntFlag): - value: Any \ No newline at end of file + + value: Any + + def __init__(self) -> None: pass From b3653d7d0b1ac54d434effa277b955e3e9aa19cd Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 09:24:04 -0400 Subject: [PATCH 06/13] Remove some unused fixtures. --- test-data/unit/check-enum.test | 4 ---- 1 file changed, 4 deletions(-) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index e4e9e4bee0bd..e3c233f32697 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -103,7 +103,6 @@ def infer_truth(truth: Truth) -> None: [case testEnumValueCustomAuto] from enum import Enum, auto -from typing import List, Any class AutoName(Enum): # In `typeshed`, this is a staticmethod and has more arguments, @@ -117,9 +116,6 @@ class Truth(AutoName): def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.str' -[builtins fixtures/bool.pyi] -[builtins fixtures/staticmethod.pyi] -[builtins fixtures/list.pyi] [case testEnumValueInhomogenous] from enum import Enum From f1adaf524d5588237916d556000bb67cd1dd363a Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 11:08:51 -0400 Subject: [PATCH 07/13] Make the self-test happy. --- mypy/plugins/enums.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index ef513301739b..90f47e4c5d87 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -65,6 +65,8 @@ def _infer_value_type_with_auto_fallback( return None if (isinstance(proper_type, Instance) and proper_type.type.fullname == 'enum.auto'): + if not isinstance(ctx.type, Instance): + raise ValueError("An incorrect ctx.type was passed.") info = ctx.type.type # Find the first _generate_next_value_ on the mro. We need to know # if it is `Enum` because `Enum` types say that the return-value of @@ -83,7 +85,7 @@ def _infer_value_type_with_auto_fallback( return ctx.default_attr_type # This should be a `CallableType` - node_type = stnode.type + node_type = get_proper_type(stnode.type) if isinstance(node_type, CallableType): if type_with_generate_next_value.fullname == 'enum.Enum': int_type = ctx.api.named_generic_type('builtins.int', []) @@ -128,14 +130,13 @@ class SomeEnum: stnodes = (info.get(name) for name in info.names) # Enums _can_ have methods. # Omit methods for our value inference. - stnodes_non_method = ( - n for n in stnodes if not isinstance(n.type, CallableType)) node_types = ( get_proper_type(n.type) if n else None - for n in stnodes_non_method) + for n in stnodes) proper_types = ( _infer_value_type_with_auto_fallback(ctx, t) - for t in node_types) + for t in node_types + if t is None or not isinstance(t, CallableType)) underlying_type = next(proper_types, None) if underlying_type is None: return ctx.default_attr_type @@ -153,13 +154,12 @@ class SomeEnum: if stnode is None: return ctx.default_attr_type - underlying_type = get_proper_type(stnode.type) + underlying_type = _infer_value_type_with_auto_fallback( + ctx, get_proper_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 - return _infer_value_type_with_auto_fallback(ctx, underlying_type) + return underlying_type def _extract_underlying_field_name(typ: Type) -> Optional[str]: From 22c1474ea9ca6fc3e4123b2b1524a4ef3c3f345e Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 20:45:32 -0400 Subject: [PATCH 08/13] Invert logic for inferring with value-type in presence of auto. --- mypy/plugins/enums.py | 58 +++++++++++++++++++++---------------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 90f47e4c5d87..d1ab70f04bad 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -63,36 +63,36 @@ def _infer_value_type_with_auto_fallback( """ if proper_type is None: return None - if (isinstance(proper_type, Instance) and - proper_type.type.fullname == 'enum.auto'): - if not isinstance(ctx.type, Instance): - raise ValueError("An incorrect ctx.type was passed.") - info = ctx.type.type - # Find the first _generate_next_value_ on the mro. We need to know - # if it is `Enum` because `Enum` types say that the return-value of - # `_generate_next_value_` is `Any`. In reality the default `auto()` - # returns an `int` (presumably the `Any` in typeshed is to make it - # easier to subclass and change the returned type). - type_with_generate_next_value = next( - (type_info for type_info in info.mro - if type_info.names.get('_generate_next_value_')), - None) - if type_with_generate_next_value is None: - return ctx.default_attr_type - - stnode = type_with_generate_next_value.get('_generate_next_value_') - if stnode is None: - return ctx.default_attr_type - - # This should be a `CallableType` - node_type = get_proper_type(stnode.type) - if isinstance(node_type, CallableType): - if type_with_generate_next_value.fullname == 'enum.Enum': - int_type = ctx.api.named_generic_type('builtins.int', []) - return int_type - return get_proper_type(node_type.ret_type) + if not (isinstance(proper_type, Instance) or + proper_type.type.fullname != 'enum.auto'): + return proper_type + if not isinstance(ctx.type, Instance): + raise ValueError("An incorrect ctx.type was passed.") + info = ctx.type.type + # Find the first _generate_next_value_ on the mro. We need to know + # if it is `Enum` because `Enum` types say that the return-value of + # `_generate_next_value_` is `Any`. In reality the default `auto()` + # returns an `int` (presumably the `Any` in typeshed is to make it + # easier to subclass and change the returned type). + type_with_generate_next_value = next( + (type_info for type_info in info.mro + if type_info.names.get('_generate_next_value_')), + None) + if type_with_generate_next_value is None: return ctx.default_attr_type - return proper_type + + stnode = type_with_generate_next_value.get('_generate_next_value_') + if stnode is None: + return ctx.default_attr_type + + # This should be a `CallableType` + node_type = get_proper_type(stnode.type) + if isinstance(node_type, CallableType): + if type_with_generate_next_value.fullname == 'enum.Enum': + int_type = ctx.api.named_generic_type('builtins.int', []) + return int_type + return get_proper_type(node_type.ret_type) + return ctx.default_attr_type def enum_value_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: From aae9de37a29e44d1e5dd965a6bfcd223624bf45e Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 20:46:13 -0400 Subject: [PATCH 09/13] Change a ValueError to an `assert` for type collapsing. --- mypy/plugins/enums.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index d1ab70f04bad..9dfbe8dad453 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -66,8 +66,7 @@ def _infer_value_type_with_auto_fallback( if not (isinstance(proper_type, Instance) or proper_type.type.fullname != 'enum.auto'): return proper_type - if not isinstance(ctx.type, Instance): - raise ValueError("An incorrect ctx.type was passed.") + assert isinstance(ctx.type, Instance), 'An incorrect ctx.type was passed.' info = ctx.type.type # Find the first _generate_next_value_ on the mro. We need to know # if it is `Enum` because `Enum` types say that the return-value of From d392d0567cd5124fe9bf42fa7f19ca2e2bcefa13 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 20:50:15 -0400 Subject: [PATCH 10/13] Shorten some names and drop a branch or two. --- mypy/plugins/enums.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 9dfbe8dad453..101d8a55ff1c 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -73,21 +73,17 @@ def _infer_value_type_with_auto_fallback( # `_generate_next_value_` is `Any`. In reality the default `auto()` # returns an `int` (presumably the `Any` in typeshed is to make it # easier to subclass and change the returned type). - type_with_generate_next_value = next( - (type_info for type_info in info.mro - if type_info.names.get('_generate_next_value_')), - None) - if type_with_generate_next_value is None: + type_with_gnv = next( + (ti for ti in info.mro if ti.names.get('_generate_next_value_')), None) + if type_with_gnv is None: return ctx.default_attr_type - stnode = type_with_generate_next_value.get('_generate_next_value_') - if stnode is None: - return ctx.default_attr_type + stnode = type_with_gnv.names['_generate_next_value_'] # This should be a `CallableType` node_type = get_proper_type(stnode.type) if isinstance(node_type, CallableType): - if type_with_generate_next_value.fullname == 'enum.Enum': + if type_with_gnv.fullname == 'enum.Enum': int_type = ctx.api.named_generic_type('builtins.int', []) return int_type return get_proper_type(node_type.ret_type) From 050fdc271ea4ca6dff575c7207d4ee6b7cc01868 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 21:43:08 -0400 Subject: [PATCH 11/13] Get my conditional right ... I think ... --- mypy/plugins/enums.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 101d8a55ff1c..823d31ab4d45 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -63,8 +63,8 @@ def _infer_value_type_with_auto_fallback( """ if proper_type is None: return None - if not (isinstance(proper_type, Instance) or - proper_type.type.fullname != 'enum.auto'): + if not ((isinstance(proper_type, Instance) and + proper_type.type.fullname == 'enum.auto')): return proper_type assert isinstance(ctx.type, Instance), 'An incorrect ctx.type was passed.' info = ctx.type.type From 5fa5a93e751d0557494c3769dee27e77c9b44bea Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Wed, 16 Sep 2020 23:22:51 -0400 Subject: [PATCH 12/13] Fix the fixture application. --- test-data/unit/check-enum.test | 2 ++ test-data/unit/lib-stub/enum.pyi | 3 --- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index e3c233f32697..37b12a0c32eb 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -87,6 +87,7 @@ class Truth(Enum): def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.int' +[builtins fixtures/primitives.pyi] [case testEnumValueExtraMethods] from enum import Enum, auto @@ -116,6 +117,7 @@ class Truth(AutoName): def infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is 'builtins.str' +[builtins fixtures/primitives.pyi] [case testEnumValueInhomogenous] from enum import Enum diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi index 29178732ea32..8d0e5fce291a 100644 --- a/test-data/unit/lib-stub/enum.pyi +++ b/test-data/unit/lib-stub/enum.pyi @@ -41,7 +41,4 @@ class IntFlag(int, Flag): class auto(IntFlag): - value: Any - - def __init__(self) -> None: pass From 5b8baa95d0e87316b6b4ab0e6b72d796767d0a26 Mon Sep 17 00:00:00 2001 From: Matt Gilson Date: Thu, 17 Sep 2020 08:54:46 -0400 Subject: [PATCH 13/13] Stop using `next` for now. --- mypy/plugins/enums.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 823d31ab4d45..e246e9de14b6 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -10,7 +10,7 @@ we actually bake some of it directly in to the semantic analysis layer (see semanal_enum.py). """ -from typing import Optional +from typing import Iterable, Optional, TypeVar from typing_extensions import Final import mypy.plugin # To avoid circular imports. @@ -53,6 +53,19 @@ def enum_name_callback(ctx: 'mypy.plugin.AttributeContext') -> Type: return str_type.copy_modified(last_known_value=literal_type) +_T = TypeVar('_T') + + +def _first(it: Iterable[_T]) -> Optional[_T]: + """Return the first value from any iterable. + + Returns ``None`` if the iterable is empty. + """ + for val in it: + return val + return None + + def _infer_value_type_with_auto_fallback( ctx: 'mypy.plugin.AttributeContext', proper_type: Optional[ProperType]) -> Optional[Type]: @@ -73,8 +86,8 @@ def _infer_value_type_with_auto_fallback( # `_generate_next_value_` is `Any`. In reality the default `auto()` # returns an `int` (presumably the `Any` in typeshed is to make it # easier to subclass and change the returned type). - type_with_gnv = next( - (ti for ti in info.mro if ti.names.get('_generate_next_value_')), None) + type_with_gnv = _first( + ti for ti in info.mro if ti.names.get('_generate_next_value_')) if type_with_gnv is None: return ctx.default_attr_type @@ -132,7 +145,7 @@ class SomeEnum: _infer_value_type_with_auto_fallback(ctx, t) for t in node_types if t is None or not isinstance(t, CallableType)) - underlying_type = next(proper_types, None) + underlying_type = _first(proper_types) if underlying_type is None: return ctx.default_attr_type all_same_value_type = all(