Add StrEnum support for Python 3.11 (#12035)

sobolevn · web-flow · commit 37d58c310a0b · 2022-02-08T22:38:14.000-08:00
diff --git a/mypy/semanal_enum.py b/mypy/semanal_enum.py
@@ -18,7 +18,7 @@
 # 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.
 ENUM_BASES: Final = frozenset((
-    'enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag',
+    'enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag', 'enum.StrEnum',
 ))
 ENUM_SPECIAL_PROPS: Final = frozenset((
     'name', 'value', '_name_', '_value_', '_order_', '__order__',
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
@@ -232,6 +232,20 @@ an_enum = SomeIntEnum.x
 an_enum = returns_some_int_enum()
 [out]
 
+[case testStrEnumCreation]
+# flags: --python-version 3.11
+from enum import StrEnum
+
+class MyStrEnum(StrEnum):
+    x = 'x'
+    y = 'y'
+
+reveal_type(MyStrEnum.x)  # N: Revealed type is "Literal[__main__.MyStrEnum.x]?"
+reveal_type(MyStrEnum.x.value)  # N: Revealed type is "Literal['x']?"
+reveal_type(MyStrEnum.y)  # N: Revealed type is "Literal[__main__.MyStrEnum.y]?"
+reveal_type(MyStrEnum.y.value)  # N: Revealed type is "Literal['y']?"
+[out]
+
 [case testEnumMethods]
 from enum import Enum
 
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
@@ -43,3 +43,8 @@ class IntFlag(int, Flag):
 
 class auto(IntFlag):
     value: Any
+
+
+# It is python-3.11+ only:
+class StrEnum(str, Enum):
+    def __new__(cls: Type[_T], value: str | _T) -> _T: ...
