Skip to content

Commit

Permalink
Fix enum with property match type exhaustion
Browse files Browse the repository at this point in the history
  • Loading branch information
terencehonles committed May 2, 2024
1 parent fb31409 commit 0f622ed
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 13 deletions.
16 changes: 4 additions & 12 deletions mypy/typeops.py
Original file line number Diff line number Diff line change
Expand Up @@ -878,18 +878,10 @@ 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
8 changes: 7 additions & 1 deletion mypy/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -1496,7 +1496,13 @@ 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
36 changes: 36 additions & 0 deletions test-data/unit/check-python310.test
Original file line number Diff line number Diff line change
Expand Up @@ -1484,6 +1484,42 @@ def g(m: Medal) -> int:
reveal_type(m) # N: Revealed type is "Literal[__main__.Medal.bronze]"
return 2


[case testMatchLiteralPatternEnumWithProperty]
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 testMatchLiteralPatternEnumCustomEquals-skip]
from enum import Enum
class Medal(Enum):
Expand Down

0 comments on commit 0f622ed

Please sign in to comment.