Makes Enum members implicitly final, refs #5599

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)

mypy/semanal.py

@@ -2408,10 +2408,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)

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