Skip to content

Enum now accepts String literals and final ... #8664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -171,6 +172,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