Enum now accepts String literals and final values as 2nd arg (#8664)

Related #8219

Co-authored-by: Jingchen Ye <[email protected]>

Author: vincent-prz

mypy/semanal_enum.py

@@ -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.
@@ -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(

test-data/unit/check-enum.test

@@ -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):