From 360fd99931414156692c2eb3eabf4d8a16feff20 Mon Sep 17 00:00:00 2001
From: Petter Friberg <petter_friberg@hotmail.com>
Date: Sun, 6 Feb 2022 21:21:25 +0100
Subject: [PATCH] Support overriding dunder attributes in Enum subclass

---
 mypy/checker.py                | 24 ++++++++++++----
 mypy/message_registry.py       |  5 ++++
 mypy/semanal.py                |  6 ++--
 test-data/unit/check-enum.test | 51 ++++++++++++++++++++++++++++++++++
 4 files changed, 79 insertions(+), 7 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index e0f2398b67fc..a40c4f4194d1 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -1834,11 +1834,7 @@ def visit_class_def(self, defn: ClassDef) -> None:
         if typ.is_protocol and typ.defn.type_vars:
             self.check_protocol_variance(defn)
         if not defn.has_incompatible_baseclass and defn.info.is_enum:
-            for base in defn.info.mro[1:-1]:  # we don't need self and `object`
-                if base.is_enum and base.fullname not in ENUM_BASES:
-                    self.check_final_enum(defn, base)
-            self.check_enum_bases(defn)
-            self.check_enum_new(defn)
+            self.check_enum(defn)
 
     def check_final_deletable(self, typ: TypeInfo) -> None:
         # These checks are only for mypyc. Only perform some checks that are easier
@@ -1896,6 +1892,24 @@ def check_init_subclass(self, defn: ClassDef) -> None:
             # all other bases have already been checked.
             break
 
+    def check_enum(self, defn: ClassDef) -> None:
+        assert defn.info.is_enum
+        if defn.info.fullname not in ENUM_BASES:
+            for sym in defn.info.names.values():
+                if (isinstance(sym.node, Var) and sym.node.has_explicit_value and
+                        sym.node.name == '__members__'):
+                    # `__members__` will always be overwritten by `Enum` and is considered
+                    # read-only so we disallow assigning a value to it
+                    self.fail(
+                        message_registry.ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN, sym.node
+                    )
+        for base in defn.info.mro[1:-1]:  # we don't need self and `object`
+            if base.is_enum and base.fullname not in ENUM_BASES:
+                self.check_final_enum(defn, base)
+
+        self.check_enum_bases(defn)
+        self.check_enum_new(defn)
+
     def check_final_enum(self, defn: ClassDef, base: TypeInfo) -> None:
         for sym in base.names.values():
             if self.is_final_enum_value(sym):
diff --git a/mypy/message_registry.py b/mypy/message_registry.py
index 68e739b24358..804fdb54be9e 100644
--- a/mypy/message_registry.py
+++ b/mypy/message_registry.py
@@ -205,6 +205,11 @@ def format(self, *args: object, **kwargs: object) -> "ErrorMessage":
 )
 CANNOT_MAKE_DELETABLE_FINAL: Final = ErrorMessage("Deletable attribute cannot be final")
 
+# Enum
+ENUM_MEMBERS_ATTR_WILL_BE_OVERRIDEN: Final = ErrorMessage(
+    'Assigned "__members__" will be overriden by "Enum" internally'
+)
+
 # ClassVar
 CANNOT_OVERRIDE_INSTANCE_VAR: Final = ErrorMessage(
     'Cannot override instance variable (previously declared on base class "{}") with class '
diff --git a/mypy/semanal.py b/mypy/semanal.py
index 60127961b60e..ef10b62086b3 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -116,6 +116,7 @@
 )
 from mypy.util import (
     correct_relative_import, unmangle, module_prefix, is_typeshed_file, unnamed_function,
+    is_dunder,
 )
 from mypy.scope import Scope
 from mypy.semanal_shared import (
@@ -2480,8 +2481,9 @@ def store_final_status(self, s: AssignmentStmt) -> None:
                     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
+                        # Double underscored members are writable on an `Enum`.
+                        # (Except read-only `__members__` but that is handled in type checker)
+                        cur_node.node.is_final = s.is_final_def = not is_dunder(cur_node.node.name)
 
                 # Special case: deferred initialization of a final attribute in __init__.
                 # In this case we just pretend this is a valid final definition to suppress
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 78460909cd6a..9e2b38bd9c8d 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -2029,3 +2029,54 @@ class C(Enum):
 
 C._ignore_ # E: "Type[C]" has no attribute "_ignore_"
 [typing fixtures/typing-medium.pyi]
+
+[case testCanOverrideDunderAttributes]
+import typing
+from enum import Enum, Flag
+
+class BaseEnum(Enum):
+    __dunder__ = 1
+    __labels__: typing.Dict[int, str]
+
+class Override(BaseEnum):
+    __dunder__ = 2
+    __labels__ = {1: "1"}
+
+Override.__dunder__ = 3
+BaseEnum.__dunder__ = 3
+Override.__labels__ = {2: "2"}
+
+class FlagBase(Flag):
+    __dunder__ = 1
+    __labels__: typing.Dict[int, str]
+
+class FlagOverride(FlagBase):
+    __dunder__ = 2
+    __labels = {1: "1"}
+
+FlagOverride.__dunder__ = 3
+FlagBase.__dunder__ = 3
+FlagOverride.__labels__ = {2: "2"}
+[builtins fixtures/dict.pyi]
+
+[case testCanNotInitialize__members__]
+import typing
+from enum import Enum
+
+class WritingMembers(Enum):
+    __members__: typing.Dict[Enum, Enum] = {}  # E: Assigned "__members__" will be overriden by "Enum" internally
+
+class OnlyAnnotatedMembers(Enum):
+    __members__: typing.Dict[Enum, Enum]
+[builtins fixtures/dict.pyi]
+
+[case testCanOverrideDunderOnNonFirstBaseEnum]
+import typing
+from enum import Enum
+
+class Some:
+    __labels__: typing.Dict[int, str]
+
+class A(Some, Enum):
+    __labels__ = {1: "1"}
+[builtins fixtures/dict.pyi]