Skip to content

Commit

Permalink
Support overriding dunder attributes in Enum subclass (#12138)
Browse files Browse the repository at this point in the history
Allows any dunder (`__name__`) attributes except `__members__` (due to it being 
read-only) to be overridden in enum subclasses.

Fixes #12132.
  • Loading branch information
Petter Friberg authored and JukkaL committed Mar 22, 2022
1 parent 837543e commit 7e09c2a
Show file tree
Hide file tree
Showing 4 changed files with 79 additions and 7 deletions.
24 changes: 19 additions & 5 deletions mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -1829,11 +1829,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
Expand Down Expand Up @@ -1891,6 +1887,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):
Expand Down
5 changes: 5 additions & 0 deletions mypy/message_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,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 '
Expand Down
6 changes: 4 additions & 2 deletions mypy/semanal.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -2473,8 +2474,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
Expand Down
51 changes: 51 additions & 0 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -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]

0 comments on commit 7e09c2a

Please sign in to comment.