[ty] Expansion of enums into unions of literals

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -2355,12 +2355,13 @@ import enum
 
 reveal_type(enum.Enum.__members__)  # revealed: MappingProxyType[str, Unknown]
 
-class Foo(enum.Enum):
-    BAR = 1
+class Answer(enum.Enum):
+    NO = 0
+    YES = 1
 
-reveal_type(Foo.BAR)  # revealed: Literal[Foo.BAR]
-reveal_type(Foo.BAR.value)  # revealed: Any
-reveal_type(Foo.__members__)  # revealed: MappingProxyType[str, Unknown]
+reveal_type(Answer.NO)  # revealed: Literal[Answer.NO]
+reveal_type(Answer.NO.value)  # revealed: Any
+reveal_type(Answer.__members__)  # revealed: MappingProxyType[str, Unknown]
 ```
 
 ## References

crates/ty_python_semantic/resources/mdtest/call/overloads.md

@@ -369,6 +369,8 @@ def _(x: type[A | B]):
 
 ### Expanding enums
 
+#### Basic
+
 `overloaded.pyi`:
 
 ```pyi
@@ -394,15 +396,106 @@ def f(x: Literal[SomeEnum.C]) -> C: ...
 ```
 
 ```py
+from typing import Literal
 from overloaded import SomeEnum, A, B, C, f
 
-def _(x: SomeEnum):
+def _(x: SomeEnum, y: Literal[SomeEnum.A, SomeEnum.C]):
     reveal_type(f(SomeEnum.A))  # revealed: A
     reveal_type(f(SomeEnum.B))  # revealed: B
     reveal_type(f(SomeEnum.C))  # revealed: C
-    # TODO: This should not be an error. The return type should be `A | B | C` once enums are expanded
-    # error: [no-matching-overload]
-    reveal_type(f(x))  # revealed: Unknown
+    reveal_type(f(x))  # revealed: A | B | C
+    reveal_type(f(y))  # revealed: A | C
+```
+
+#### Enum with single member
+
+This pattern appears in typeshed. Here, it is used to represent two optional, mutually exclusive
+keyword parameters:
+
+`overloaded.pyi`:
+
+```pyi
+from enum import Enum, auto
+from typing import overload, Literal
+
+class Missing(Enum):
+    Value = auto()
+
+class OnlyASpecified: ...
+class OnlyBSpecified: ...
+class BothMissing: ...
+
+@overload
+def f(*, a: int, b: Literal[Missing.Value] = ...) -> OnlyASpecified: ...
+@overload
+def f(*, a: Literal[Missing.Value] = ..., b: int) -> OnlyBSpecified: ...
+@overload
+def f(*, a: Literal[Missing.Value] = ..., b: Literal[Missing.Value] = ...) -> BothMissing: ...
+```
+
+```py
+from typing import Literal
+from overloaded import f, Missing
+
+reveal_type(f())  # revealed: BothMissing
+reveal_type(f(a=0))  # revealed: OnlyASpecified
+reveal_type(f(b=0))  # revealed: OnlyBSpecified
+
+f(a=0, b=0)  # error: [no-matching-overload]
+
+def _(missing: Literal[Missing.Value], missing_or_present: Literal[Missing.Value] | int):
+    reveal_type(f(a=missing, b=missing))  # revealed: BothMissing
+    reveal_type(f(a=missing))  # revealed: BothMissing
+    reveal_type(f(b=missing))  # revealed: BothMissing
+    reveal_type(f(a=0, b=missing))  # revealed: OnlyASpecified
+    reveal_type(f(a=missing, b=0))  # revealed: OnlyBSpecified
+
+    reveal_type(f(a=missing_or_present))  # revealed: BothMissing | OnlyASpecified
+    reveal_type(f(b=missing_or_present))  # revealed: BothMissing | OnlyBSpecified
+
+    # Here, both could be present, so this should be an error
+    f(a=missing_or_present, b=missing_or_present)  # error: [no-matching-overload]
+```
+
+#### Enum subclass without members
+
+An `Enum` subclass without members should *not* be expanded:
+
+`overloaded.pyi`:
+
+```pyi
+from enum import Enum
+from typing import overload, Literal
+
+class MyEnumSubclass(Enum):
+    pass
+
+class ActualEnum(MyEnumSubclass):
+    A = 1
+    B = 2
+
+class OnlyA: ...
+class OnlyB: ...
+class Both: ...
+
+@overload
+def f(x: Literal[ActualEnum.A]) -> OnlyA: ...
+@overload
+def f(x: Literal[ActualEnum.B]) -> OnlyB: ...
+@overload
+def f(x: ActualEnum) -> Both: ...
+@overload
+def f(x: MyEnumSubclass) -> MyEnumSubclass: ...
+```
+
+```py
+from overloaded import MyEnumSubclass, ActualEnum, f
+
+def _(actual_enum: ActualEnum, my_enum_instance: MyEnumSubclass):
+    reveal_type(f(actual_enum))  # revealed: Both
+    reveal_type(f(ActualEnum.A))  # revealed: OnlyA
+    reveal_type(f(ActualEnum.B))  # revealed: OnlyB
+    reveal_type(f(my_enum_instance))  # revealed: MyEnumSubclass
 ```
 
 ### No matching overloads

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -570,7 +570,111 @@ To do: <https://typing.python.org/en/latest/spec/enums.html#enum-definition>
 
 ## Exhaustiveness checking
 
-To do
+## `if` statements
+
+```py
+from enum import Enum
+from typing_extensions import assert_never
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+def color_name(color: Color) -> str:
+    if color is Color.RED:
+        return "Red"
+    elif color is Color.GREEN:
+        return "Green"
+    elif color is Color.BLUE:
+        return "Blue"
+    else:
+        assert_never(color)
+
+# No `invalid-return-type` error here because the implicit `else` branch is detected as unreachable:
+def color_name_without_assertion(color: Color) -> str:
+    if color is Color.RED:
+        return "Red"
+    elif color is Color.GREEN:
+        return "Green"
+    elif color is Color.BLUE:
+        return "Blue"
+
+def color_name_misses_one_variant(color: Color) -> str:
+    if color is Color.RED:
+        return "Red"
+    elif color is Color.GREEN:
+        return "Green"
+    else:
+        assert_never(color)  # error: [type-assertion-failure] "Argument does not have asserted type `Never`"
+
+class Singleton(Enum):
+    VALUE = 1
+
+def singleton_check(value: Singleton) -> str:
+    if value is Singleton.VALUE:
+        return "Singleton value"
+    else:
+        assert_never(value)
+```
+
+## `match` statements
+
+```toml
+[environment]
+python-version = "3.10"
+```
+
+```py
+from enum import Enum
+from typing_extensions import assert_never
+
+class Color(Enum):
+    RED = 1
+    GREEN = 2
+    BLUE = 3
+
+def color_name(color: Color) -> str:
+    match color:
+        case Color.RED:
+            return "Red"
+        case Color.GREEN:
+            return "Green"
+        case Color.BLUE:
+            return "Blue"
+        case _:
+            assert_never(color)
+
+# TODO: this should not be an error, see https://github.com/astral-sh/ty/issues/99#issuecomment-2983054488
+# error: [invalid-return-type] "Function can implicitly return `None`, which is not assignable to return type `str`"
+def color_name_without_assertion(color: Color) -> str:
+    match color:
+        case Color.RED:
+            return "Red"
+        case Color.GREEN:
+            return "Green"
+        case Color.BLUE:
+            return "Blue"
+
+def color_name_misses_one_variant(color: Color) -> str:
+    match color:
+        case Color.RED:
+            return "Red"
+        case Color.GREEN:
+            return "Green"
+        case _:
+            assert_never(color)  # error: [type-assertion-failure] "Argument does not have asserted type `Never`"
+
+class Singleton(Enum):
+    VALUE = 1
+
+def singleton_check(value: Singleton) -> str:
+    match value:
+        case Singleton.VALUE:
+            return "Singleton value"
+        case _:
+            assert_never(value)
+```
 
 ## References
 

crates/ty_python_semantic/resources/mdtest/intersection_types.md

@@ -763,6 +763,65 @@ def f(
     reveal_type(j)  # revealed: Unknown & Literal[""]
 ```
 
+## Simplifications involving enums and enum literals
+
+```toml
+[environment]
+python-version = "3.12"
+```
+
+```py
+from ty_extensions import Intersection, Not
+from typing import Literal
+from enum import Enum
+
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+    BLUE = "blue"
+
+type Red = Literal[Color.RED]
+type Green = Literal[Color.GREEN]
+type Blue = Literal[Color.BLUE]
+
+def f(
+    a: Intersection[Color, Red],
+    b: Intersection[Color, Not[Red]],
+    c: Intersection[Color, Not[Red | Green]],
+    d: Intersection[Color, Not[Red | Green | Blue]],
+    e: Intersection[Red, Not[Color]],
+    f: Intersection[Red | Green, Not[Color]],
+    g: Intersection[Not[Red], Color],
+    h: Intersection[Red, Green],
+    i: Intersection[Red | Green, Green | Blue],
+):
+    reveal_type(a)  # revealed: Literal[Color.RED]
+    reveal_type(b)  # revealed: Literal[Color.GREEN, Color.BLUE]
+    reveal_type(c)  # revealed: Literal[Color.BLUE]
+    reveal_type(d)  # revealed: Never
+    reveal_type(e)  # revealed: Never
+    reveal_type(f)  # revealed: Never
+    reveal_type(g)  # revealed: Literal[Color.GREEN, Color.BLUE]
+    reveal_type(h)  # revealed: Never
+    reveal_type(i)  # revealed: Literal[Color.GREEN]
+
+class Single(Enum):
+    VALUE = 0
+
+def g(
+    a: Intersection[Single, Literal[Single.VALUE]],
+    b: Intersection[Single, Not[Literal[Single.VALUE]]],
+    c: Intersection[Not[Literal[Single.VALUE]], Single],
+    d: Intersection[Single, Not[Single]],
+    e: Intersection[Single | int, Not[Single]],
+):
+    reveal_type(a)  # revealed: Single
+    reveal_type(b)  # revealed: Never
+    reveal_type(c)  # revealed: Never
+    reveal_type(d)  # revealed: Never
+    reveal_type(e)  # revealed: int
+```
+
 ## Addition of a type to an intersection with many non-disjoint types
 
 This slightly strange-looking test is a regression test for a mistake that was nearly made in a PR:

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md

@@ -1,4 +1,4 @@
-# Narrowing for `!=` conditionals
+# Narrowing for `!=` and `==` conditionals
 
 ## `x != None`
 
@@ -22,6 +22,12 @@ def _(x: bool):
         reveal_type(x)  # revealed: Literal[True]
     else:
         reveal_type(x)  # revealed: Literal[False]
+
+def _(x: bool):
+    if x == False:
+        reveal_type(x)  # revealed: Literal[False]
+    else:
+        reveal_type(x)  # revealed: Literal[True]
 ```
 
 ### Enums
@@ -35,11 +41,31 @@ class Answer(Enum):
 
 def _(answer: Answer):
     if answer != Answer.NO:
-        # TODO: This should be simplified to `Literal[Answer.YES]`
-        reveal_type(answer)  # revealed: Answer & ~Literal[Answer.NO]
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
+    else:
+        reveal_type(answer)  # revealed: Literal[Answer.NO]
+
+def _(answer: Answer):
+    if answer == Answer.NO:
+        reveal_type(answer)  # revealed: Literal[Answer.NO]
+    else:
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
+
+class Single(Enum):
+    VALUE = 1
+
+def _(x: Single | int):
+    if x != Single.VALUE:
+        reveal_type(x)  # revealed: int
+    else:
+        # `int` is not eliminated here because there could be subclasses of `int` with custom `__eq__`/`__ne__` methods
+        reveal_type(x)  # revealed: Single | int
+
+def _(x: Single | int):
+    if x == Single.VALUE:
+        reveal_type(x)  # revealed: Single | int
     else:
-        # TODO: This should be `Literal[Answer.NO]`
-        reveal_type(answer)  # revealed: Answer
+        reveal_type(x)  # revealed: int
 ```
 
 This narrowing behavior is only safe if the enum has no custom `__eq__`/`__ne__` method:

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/is.md

@@ -78,8 +78,16 @@ def _(answer: Answer):
     if answer is Answer.NO:
         reveal_type(answer)  # revealed: Literal[Answer.NO]
     else:
-        # TODO: This should be `Literal[Answer.YES]`
-        reveal_type(answer)  # revealed: Answer & ~Literal[Answer.NO]
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
+
+class Single(Enum):
+    VALUE = 1
+
+def _(x: Single | int):
+    if x is Single.VALUE:
+        reveal_type(x)  # revealed: Single
+    else:
+        reveal_type(x)  # revealed: int
 ```
 
 ## `is` for `EllipsisType` (Python 3.10+)