diff --git a/mypy/plugins/enums.py b/mypy/plugins/enums.py index 1b22c09fe7bb5..53241304cf48a 100644 --- a/mypy/plugins/enums.py +++ b/mypy/plugins/enums.py @@ -10,12 +10,14 @@ 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. 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,44 @@ 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 + + # 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 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 + + # But, after we started treating all `Enum` values as `Final`, + # 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: + # + # 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) + if all_equivalent_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 2748c7b17389c..3c0d548b2747e 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py @@ -2411,10 +2411,30 @@ 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 22d167f3487b4..21350c0301865 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 @@ -1362,6 +1376,25 @@ class E(IntEnum): 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" + + +[case testEnumFinalValuesCannotRedefineValueProp] +from enum import Enum +class Types(Enum): + key = 0 + value = 1 # E: Cannot override writable attribute "value" with a final one + + [case testEnumReusedKeys] # https://github.com/python/mypy/issues/11248 from enum import Enum @@ -1405,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