From 1e5d993e7c9c191afe064c7e31ea6a06c7f92305 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Wed, 21 Jul 2021 20:31:16 +0300 Subject: [PATCH 1/5] Makes Enum members implicitly final, refs #5599 --- mypy/plugins/enums.py | 11 ++++++++++- mypy/semanal.py | 24 +++++++++++++++++++++--- test-data/unit/check-enum.test | 34 ++++++++++++++++++++++++++++++---- 3 files changed, 61 insertions(+), 8 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 1b22c09fe7bb..cbee28b9717c 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -15,7 +15,9 @@ import mypy.plugin # To avoid circular imports. from mypy.types import Type, Instance, LiteralType, CallableType, ProperType, get_proper_type +from mypy.typeops import make_simplified_union from mypy.nodes import TypeInfo +from mypy.subtypes import is_equivalent # 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. @@ -165,19 +167,26 @@ class SomeEnum: get_proper_type(n.type) if n else None for n in stnodes if n is None or not n.implicit) - proper_types = ( + proper_types = list( _infer_value_type_with_auto_fallback(ctx, t) for t in node_types if t is None or not isinstance(t, CallableType)) underlying_type = _first(proper_types) if underlying_type is None: return ctx.default_attr_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 + + all_equivalent_types = all( + proper_type is not None and is_equivalent(proper_type, underlying_type) + for proper_type in proper_types) + if all_equivalent_types: + return make_simplified_union(proper_types) return ctx.default_attr_type assert isinstance(ctx.type, Instance) diff --git a/mypy/semanal.py b/mypy/semanal.py index 15036e065265..86413829f0e8 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -2381,10 +2381,28 @@ def store_final_status(self, s: AssignmentStmt) -> None: (isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs)): node.final_unset_in_class = True else: - # Special case: deferred initialization of a final attribute in __init__. - # In this case we just pretend this is a valid final definition to suppress - # errors about assigning to final attribute. for lval in self.flatten_lvalues(s.lvalues): + # Special case: we are working with an `Enum`: + # + # class MyEnum(Enum): + # key = 'some value' + # + # Here `key` is implicitly final. In runtime, code like + # + # MyEnum.key = 'modified' + # + # will fail with `AttributeError: Cannot reassign members.` + # That's why we need to replicate this. + if isinstance(lval, NameExpr) and isinstance(self.type, TypeInfo) and self.type.is_enum: + cur_node = self.type.names.get(lval.name, None) + if (cur_node and isinstance(cur_node.node, Var) and + not (isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs)): + cur_node.node.is_final = True + s.is_final_def = True + + # Special case: deferred initialization of a final attribute in __init__. + # In this case we just pretend this is a valid final definition to suppress + # errors about assigning to final attribute. if isinstance(lval, MemberExpr) and self.is_self_member_ref(lval): assert self.type, "Self member outside a class" cur_node = self.type.names.get(lval.name, None) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 5200c00d3f28..d2364ae329f0 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -56,7 +56,7 @@ class Truth(Enum): x = '' x = Truth.true.name reveal_type(Truth.true.name) # N: Revealed type is "Literal['true']?" -reveal_type(Truth.false.value) # N: Revealed type is "builtins.bool" +reveal_type(Truth.false.value) # N: Revealed type is "Literal[False]?" [builtins fixtures/bool.pyi] [case testEnumValueExtended] @@ -66,7 +66,7 @@ class Truth(Enum): false = False def infer_truth(truth: Truth) -> None: - reveal_type(truth.value) # N: Revealed type is "builtins.bool" + reveal_type(truth.value) # N: Revealed type is "Union[Literal[True]?, Literal[False]?]" [builtins fixtures/bool.pyi] [case testEnumValueAllAuto] @@ -90,7 +90,7 @@ def infer_truth(truth: Truth) -> None: [builtins fixtures/primitives.pyi] [case testEnumValueExtraMethods] -from enum import Enum, auto +from enum import Enum class Truth(Enum): true = True false = False @@ -99,7 +99,7 @@ class Truth(Enum): return 'bar' def infer_truth(truth: Truth) -> None: - reveal_type(truth.value) # N: Revealed type is "builtins.bool" + reveal_type(truth.value) # N: Revealed type is "Union[Literal[True]?, Literal[False]?]" [builtins fixtures/bool.pyi] [case testEnumValueCustomAuto] @@ -129,6 +129,20 @@ def cannot_infer_truth(truth: Truth) -> None: reveal_type(truth.value) # N: Revealed type is "Any" [builtins fixtures/bool.pyi] +[case testEnumValueSameType] +from enum import Enum + +def newbool() -> bool: + ... + +class Truth(Enum): + true = newbool() + false = newbool() + +def infer_truth(truth: Truth) -> None: + reveal_type(truth.value) # N: Revealed type is "builtins.bool" +[builtins fixtures/bool.pyi] + [case testEnumUnique] import enum @enum.unique @@ -1360,3 +1374,15 @@ class E(IntEnum): A = N(0) reveal_type(E.A.value) # N: Revealed type is "__main__.N" + + +[case testEnumFinalValues] +from enum import Enum +class Medal(Enum): + gold = 1 + silver = 2 + +# Another value: +Medal.gold = 0 # E: Cannot assign to final attribute "gold" +# Same value: +Medal.silver = 2 # E: Cannot assign to final attribute "silver" From bade480ca26ca745c801ff16ff46fea811d8374f Mon Sep 17 00:00:00 2001 From: sobolevn Date: Wed, 21 Jul 2021 21:17:41 +0300 Subject: [PATCH 2/5] Fixes CI --- mypy/plugins/enums.py | 4 ++-- mypy/semanal.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index cbee28b9717c..44b1801b962f 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 Iterable, Optional, TypeVar +from typing import Iterable, Optional, Sequence, TypeVar, cast from typing_extensions import Final import mypy.plugin # To avoid circular imports. @@ -186,7 +186,7 @@ class SomeEnum: proper_type is not None and is_equivalent(proper_type, underlying_type) for proper_type in proper_types) if all_equivalent_types: - return make_simplified_union(proper_types) + return make_simplified_union(cast(Sequence[Type], proper_types)) return ctx.default_attr_type assert isinstance(ctx.type, Instance) diff --git a/mypy/semanal.py b/mypy/semanal.py index 86413829f0e8..45e3366a1a8c 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -2393,7 +2393,9 @@ def store_final_status(self, s: AssignmentStmt) -> None: # # will fail with `AttributeError: Cannot reassign members.` # That's why we need to replicate this. - if isinstance(lval, NameExpr) and isinstance(self.type, TypeInfo) and self.type.is_enum: + if (isinstance(lval, NameExpr) and + isinstance(self.type, TypeInfo) and + self.type.is_enum): cur_node = self.type.names.get(lval.name, None) if (cur_node and isinstance(cur_node.node, Var) and not (isinstance(s.rvalue, TempNode) and s.rvalue.no_rhs)): From f217224d9b6e7653ee046f0e8d8788f76b21b859 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Sat, 25 Sep 2021 19:14:05 +0300 Subject: [PATCH 3/5] Adds more tests, adds comments to `enum` plugin --- mypy/plugins/enums.py | 18 ++++++++++++++++++ test-data/unit/check-enum.test | 7 +++++++ 2 files changed, 25 insertions(+) diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 44b1801b962f..289d68b81169 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -175,6 +175,10 @@ class SomeEnum: if underlying_type is None: return ctx.default_attr_type + # At first we try to predice future `value` type if all other items + # have the same type. For example, `int`. + # If this is the case, we simply return this type. + # See https://github.com/python/mypy/pull/9443 all_same_value_type = all( proper_type is not None and proper_type == underlying_type for proper_type in proper_types) @@ -182,6 +186,20 @@ class SomeEnum: if underlying_type is not None: return underlying_type + # But, after we started treating all `Enum` values as `Final`, + # we start to inference types in + # `item = 1` as `Literal[1]`, not just `int`. + # So, for example types in this `Enum` will all be different: + # + # class Ordering(IntEnum): + # one = 1 + # two = 2 + # three = 3 + # + # We will infer three `Literal` types here. + # They are not the same, but they are equivalent. + # So, we unify them to make sure `.value` prediction still works. + # Result will be `Literal[1] | Literal[2] | Literal[3]` for this case. all_equivalent_types = all( proper_type is not None and is_equivalent(proper_type, underlying_type) for proper_type in proper_types) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index d2364ae329f0..bce84127a25b 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -1386,3 +1386,10 @@ class Medal(Enum): Medal.gold = 0 # E: Cannot assign to final attribute "gold" # Same value: Medal.silver = 2 # E: Cannot assign to final attribute "silver" + + +[case testEnumFinalValuesCannotRedefineValueProp] +from enum import Enum +class Types(Enum): + key = 0 + value = 1 # E: Cannot override writable attribute "value" with a final one From 5e8d722f9a9601833d11ed47530ad15f988dabf3 Mon Sep 17 00:00:00 2001 From: Shantanu <12621235+hauntsaninja@users.noreply.github.com> Date: Sun, 3 Oct 2021 15:03:56 -0700 Subject: [PATCH 4/5] Typos --- 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 289d68b81169..53241304cf48 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -175,7 +175,7 @@ class SomeEnum: if underlying_type is None: return ctx.default_attr_type - # At first we try to predice future `value` type if all other items + # At first we try to predict future `value` type if all other items # have the same type. For example, `int`. # If this is the case, we simply return this type. # See https://github.com/python/mypy/pull/9443 @@ -187,7 +187,7 @@ class SomeEnum: return underlying_type # But, after we started treating all `Enum` values as `Final`, - # we start to inference types in + # we start to infer types in # `item = 1` as `Literal[1]`, not just `int`. # So, for example types in this `Enum` will all be different: # From 8b8c6554805baacfe55ac2fa9630e547fa85833e Mon Sep 17 00:00:00 2001 From: sobolevn Date: Sun, 7 Nov 2021 11:22:33 +0300 Subject: [PATCH 5/5] Fixes tests --- test-data/unit/check-enum.test | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index beaabbb4f699..21350c030186 100644 --- a/test-data/unit/check-enum.test +++ b/test-data/unit/check-enum.test @@ -1438,13 +1438,13 @@ class NonEmptyIntFlag(IntFlag): x = 1 class ErrorEnumWithValue(NonEmptyEnum): # E: Cannot inherit from final class "NonEmptyEnum" - x = 1 + x = 1 # E: Cannot override final attribute "x" (previously declared in base class "NonEmptyEnum") class ErrorIntEnumWithValue(NonEmptyIntEnum): # E: Cannot inherit from final class "NonEmptyIntEnum" - x = 1 + x = 1 # E: Cannot override final attribute "x" (previously declared in base class "NonEmptyIntEnum") class ErrorFlagWithValue(NonEmptyFlag): # E: Cannot inherit from final class "NonEmptyFlag" - x = 1 + x = 1 # E: Cannot override final attribute "x" (previously declared in base class "NonEmptyFlag") class ErrorIntFlagWithValue(NonEmptyIntFlag): # E: Cannot inherit from final class "NonEmptyIntFlag" - x = 1 + x = 1 # E: Cannot override final attribute "x" (previously declared in base class "NonEmptyIntFlag") class ErrorEnumWithoutValue(NonEmptyEnum): # E: Cannot inherit from final class "NonEmptyEnum" pass