Skip to content

Commit

Permalink
Fix enum attributes are not members
Browse files Browse the repository at this point in the history
This adds on to the change in #17182
and fixes enum attributes being used as members.
  • Loading branch information
terencehonles committed Jun 5, 2024
1 parent fbaa7e0 commit 8eb4218
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 14 deletions.
6 changes: 6 additions & 0 deletions mypy/semanal_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,12 @@ def build_enum_call_typeinfo(
var = Var(item)
var.info = info
var.is_property = True
# When an enum is created by its functional form `Enum(name, values)`
# - if it is a string it is first split by commas/whitespace
# - if it is an iterable of single items each item is assigned a value starting at `start`
# - if it is an iterable of (name, value) then the given values will be used
# either way, each item should be treated as if it has an explicit value.
var.has_explicit_value = True
var._fullname = f"{info.fullname}.{item}"
info.names[item] = SymbolTableNode(MDEF, var)
return info
Expand Down
16 changes: 3 additions & 13 deletions mypy/typeops.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@
)
from mypy.state import state
from mypy.types import (
ENUM_REMOVED_PROPS,
AnyType,
CallableType,
ExtraAttrs,
Expand Down Expand Up @@ -878,18 +877,9 @@ class Status(Enum):
return make_simplified_union(items, contract_literals=False)
elif isinstance(typ, Instance) and typ.type.fullname == target_fullname:
if typ.type.is_enum:
new_items = []
for name, symbol in typ.type.names.items():
if not isinstance(symbol.node, Var):
continue
# Skip these since Enum will remove it
if name in ENUM_REMOVED_PROPS:
continue
# Skip private attributes
if name.startswith("__"):
continue
new_items.append(LiteralType(name, typ))
return make_simplified_union(new_items, contract_literals=False)
return make_simplified_union(
[LiteralType(name, typ) for name in typ.get_enum_values()], contract_literals=False
)
elif typ.type.fullname == "builtins.bool":
return make_simplified_union(
[LiteralType(True, typ), LiteralType(False, typ)], contract_literals=False
Expand Down
9 changes: 8 additions & 1 deletion mypy/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -1512,7 +1512,14 @@ def is_singleton_type(self) -> bool:
def get_enum_values(self) -> list[str]:
"""Return the list of values for an Enum."""
return [
name for name, sym in self.type.names.items() if isinstance(sym.node, mypy.nodes.Var)
name
for name, sym in self.type.names.items()
if (
isinstance(sym.node, mypy.nodes.Var)
and name not in ENUM_REMOVED_PROPS
and not name.startswith("__")
and sym.node.has_explicit_value
)
]


Expand Down
54 changes: 54 additions & 0 deletions test-data/unit/check-python310.test
Original file line number Diff line number Diff line change
Expand Up @@ -1499,6 +1499,60 @@ def g(m: Medal) -> int:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.bronze]"
return 2


[case testMatchLiteralPatternEnumWithTypedAttribute]
from enum import Enum
from typing import NoReturn
def assert_never(x: NoReturn) -> None: ...

class int:
def __new__(cls, value: int): pass

class Medal(int, Enum):
prize: str

def __new__(cls, value: int, prize: str) -> Medal:
enum = int.__new__(cls, value)
enum._value_ = value
enum.prize = prize
return enum

gold = (1, 'cash prize')
silver = (2, 'sponsorship')
bronze = (3, 'nothing')

m: Medal

match m:
case Medal.gold:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.gold]"
case Medal.silver:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.silver]"
case Medal.bronze:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.bronze]"
case _ as unreachable:
assert_never(unreachable)

[builtins fixtures/tuple.pyi]

[case testMatchLiteralPatternFunctionalEnum]
from enum import Enum
from typing import NoReturn
def assert_never(x: NoReturn) -> None: ...

Medal = Enum('Medal', 'gold silver bronze')
m: Medal

match m:
case Medal.gold:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.gold]"
case Medal.silver:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.silver]"
case Medal.bronze:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.bronze]"
case _ as unreachable:
assert_never(unreachable)

[case testMatchLiteralPatternEnumCustomEquals-skip]
from enum import Enum
class Medal(Enum):
Expand Down

0 comments on commit 8eb4218

Please sign in to comment.