Skip to content

Commit

Permalink
Enum now accepts String literals and final values as 2nd arg (#8664)
Browse files Browse the repository at this point in the history
Related #8219

Co-authored-by: Jingchen Ye <97littleleaf11@gmail.com>
  • Loading branch information
vincent-prz and 97littleleaf11 authored Jan 7, 2022
1 parent 10f6e4d commit f5ac47f
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 0 deletions.
18 changes: 18 additions & 0 deletions mypy/semanal_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
)
from mypy.semanal_shared import SemanticAnalyzerInterface
from mypy.options import Options
from mypy.types import get_proper_type, LiteralType

# 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.
Expand Down Expand Up @@ -173,6 +174,23 @@ def parse_enum_call_args(self, call: CallExpr,
"%s() with dict literal requires string literals" % class_name, call)
items.append(key.value)
values.append(value)
elif isinstance(args[1], RefExpr) and isinstance(args[1].node, Var):
proper_type = get_proper_type(args[1].node.type)
if (proper_type is not None
and isinstance(proper_type, LiteralType)
and isinstance(proper_type.value, str)):
fields = proper_type.value
for field in fields.replace(',', ' ').split():
items.append(field)
elif args[1].node.is_final and isinstance(args[1].node.final_value, str):
fields = args[1].node.final_value
for field in fields.replace(',', ' ').split():
items.append(field)
else:
return self.fail_enum_call_arg(
"%s() expects a string, tuple, list or dict literal as the second argument" %
class_name,
call)
else:
# TODO: Allow dict(x=1, y=2) as a substitute for {'x': 1, 'y': 2}?
return self.fail_enum_call_arg(
Expand Down
32 changes: 32 additions & 0 deletions test-data/unit/check-enum.test
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,38 @@ m = Medal.gold
if int():
m = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "Medal")

-- Creation from Enum call
-- -----------------------

[case testEnumCreatedFromStringLiteral]
from enum import Enum
from typing_extensions import Literal

x: Literal['ANT BEE CAT DOG'] = 'ANT BEE CAT DOG'
Animal = Enum('Animal', x)
reveal_type(Animal.ANT) # N: Revealed type is "Literal[__main__.Animal.ANT]?"
reveal_type(Animal.BEE) # N: Revealed type is "Literal[__main__.Animal.BEE]?"
reveal_type(Animal.CAT) # N: Revealed type is "Literal[__main__.Animal.CAT]?"
reveal_type(Animal.DOG) # N: Revealed type is "Literal[__main__.Animal.DOG]?"

[builtins fixtures/tuple.pyi]

[case testEnumCreatedFromFinalValue]
from enum import Enum
from typing_extensions import Final

x: Final['str'] = 'ANT BEE CAT DOG'
Animal = Enum('Animal', x)
reveal_type(Animal.ANT) # N: Revealed type is "Literal[__main__.Animal.ANT]?"
reveal_type(Animal.BEE) # N: Revealed type is "Literal[__main__.Animal.BEE]?"
reveal_type(Animal.CAT) # N: Revealed type is "Literal[__main__.Animal.CAT]?"
reveal_type(Animal.DOG) # N: Revealed type is "Literal[__main__.Animal.DOG]?"

[builtins fixtures/tuple.pyi]

-- Creation from EnumMeta
-- ----------------------

[case testEnumFromEnumMetaBasics]
from enum import EnumMeta
class Medal(metaclass=EnumMeta):
Expand Down

0 comments on commit f5ac47f

Please sign in to comment.