[ty] Expansion of enums into unions of literals (#19382)

## Summary

Implement expansion of enums into unions of enum literals (and the
reverse operation). For the enum below, this allows us to understand
that `Color = Literal[Color.RED, Color.GREEN, Color.BLUE]`, or that
`Color & ~Literal[Color.RED] = Literal[Color.GREEN, Color.BLUE]`. This
helps in exhaustiveness checking, which is why we see some removed
`assert_never` false positives. And since exhaustiveness checking also
helps with understanding terminal control flow, we also see a few
removed `invalid-return-type` and `possibly-unresolved-reference` false
positives. This PR also adds expansion of enums in overload resolution
and type narrowing constructs.

```py
from enum import Enum
from typing_extensions import Literal, assert_never
from ty_extensions import Intersection, Not, static_assert, is_equivalent_to

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

type Red = Literal[Color.RED]
type Green = Literal[Color.GREEN]
type Blue = Literal[Color.BLUE]

static_assert(is_equivalent_to(Red | Green | Blue, Color))
static_assert(is_equivalent_to(Intersection[Color, Not[Red]], Green | Blue))


def color_name(color: Color) -> str:  # no error here (we detect that this can not implicitly return None)
    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 error here
```

## Performance

I avoided an initial regression here for large enums, but the
`UnionBuilder` and `IntersectionBuilder` parts can certainly still be
optimized. We might want to use the same technique that we also use for
unions of other literals. I didn't see any problems in our benchmarks so
far, so this is not included yet.

## Test Plan

Many new Markdown tests

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:
-        # TODO: This should be `Literal[Answer.NO]`
-        reveal_type(answer)  # revealed: Answer
+        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:
+        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+)

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

@@ -18,6 +18,8 @@ def _(flag: bool):
 
 ## `is not` for other singleton types
 
+Boolean literals:
+
 ```py
 def _(flag: bool):
     x = True if flag else False
@@ -29,6 +31,33 @@ def _(flag: bool):
         reveal_type(x)  # revealed: Literal[False]
 ```
 
+Enum literals:
+
+```py
+from enum import Enum
+
+class Answer(Enum):
+    NO = 0
+    YES = 1
+
+def _(answer: Answer):
+    if answer is not Answer.NO:
+        reveal_type(answer)  # revealed: Literal[Answer.YES]
+    else:
+        reveal_type(answer)  # revealed: Literal[Answer.NO]
+
+    reveal_type(answer)  # revealed: Answer
+
+class Single(Enum):
+    VALUE = 1
+
+def _(x: Single | int):
+    if x is not Single.VALUE:
+        reveal_type(x)  # revealed: int
+    else:
+        reveal_type(x)  # revealed: Single
+```
+
 ## `is not` for non-singleton types
 
 Non-singleton types should *not* narrow the type: two instances of a non-singleton class may occupy

crates/ty_python_semantic/resources/mdtest/type_properties/is_assignable_to.md

@@ -138,11 +138,15 @@ class Answer(Enum):
 static_assert(is_assignable_to(Literal[Answer.YES], Literal[Answer.YES]))
 static_assert(is_assignable_to(Literal[Answer.YES], Answer))
 static_assert(is_assignable_to(Literal[Answer.YES, Answer.NO], Answer))
-# TODO: this should not be an error
-# error: [static-assert-error]
 static_assert(is_assignable_to(Answer, Literal[Answer.YES, Answer.NO]))
 
 static_assert(not is_assignable_to(Literal[Answer.YES], Literal[Answer.NO]))
+
+class Single(Enum):
+    VALUE = 1
+
+static_assert(is_assignable_to(Literal[Single.VALUE], Single))
+static_assert(is_assignable_to(Single, Literal[Single.VALUE]))
 ```
 
 ### Slice literals

crates/ty_python_semantic/resources/mdtest/type_properties/is_equivalent_to.md

@@ -21,6 +21,9 @@ class Answer(Enum):
     NO = 0
     YES = 1
 
+class Single(Enum):
+    VALUE = 1
+
 static_assert(is_equivalent_to(Literal[1, 2], Literal[1, 2]))
 static_assert(is_equivalent_to(type[object], type))
 static_assert(is_equivalent_to(type, type[object]))
@@ -31,12 +34,15 @@ static_assert(not is_equivalent_to(Literal[1, 2], Literal[1, 2, 3]))
 static_assert(not is_equivalent_to(Literal[1, 2, 3], Literal[1, 2]))
 
 static_assert(is_equivalent_to(Literal[Answer.YES], Literal[Answer.YES]))
-# TODO: these should be equivalent
-# error: [static-assert-error]
+static_assert(is_equivalent_to(Literal[Answer.NO, Answer.YES], Answer))
 static_assert(is_equivalent_to(Literal[Answer.YES, Answer.NO], Answer))
 static_assert(not is_equivalent_to(Literal[Answer.YES], Literal[Answer.NO]))
 static_assert(not is_equivalent_to(Literal[Answer.YES], Answer))
 
+static_assert(is_equivalent_to(Literal[Single.VALUE], Single))
+static_assert(is_equivalent_to(Single, Literal[Single.VALUE]))
+static_assert(is_equivalent_to(Literal[Single.VALUE], Literal[Single.VALUE]))
+
 static_assert(is_equivalent_to(Never, Never))
 static_assert(is_equivalent_to(AlwaysTruthy, AlwaysTruthy))
 static_assert(is_equivalent_to(AlwaysFalsy, AlwaysFalsy))
@@ -69,8 +75,9 @@ static_assert(not is_equivalent_to(type[object], type[Any]))
 ## Unions and intersections
 
 ```py
-from typing import Any
+from typing import Any, Literal
 from ty_extensions import Intersection, Not, Unknown, is_equivalent_to, static_assert
+from enum import Enum
 
 static_assert(is_equivalent_to(str | int, str | int))
 static_assert(is_equivalent_to(str | int | Any, str | int | Unknown))
@@ -111,6 +118,11 @@ static_assert(is_equivalent_to(Intersection[P, Q], Intersection[Q, P]))
 static_assert(is_equivalent_to(Intersection[Q, Not[P]], Intersection[Not[P], Q]))
 static_assert(is_equivalent_to(Intersection[Q, R, Not[P]], Intersection[Not[P], R, Q]))
 static_assert(is_equivalent_to(Intersection[Q | R, Not[P | S]], Intersection[Not[S | P], R | Q]))
+
+class Single(Enum):
+    VALUE = 1
+
+static_assert(is_equivalent_to(P | Q | Single, Literal[Single.VALUE] | Q | P))
 ```
 
 ## Tuples

crates/ty_python_semantic/resources/mdtest/type_properties/is_single_valued.md

@@ -47,6 +47,9 @@ class NormalEnum(Enum):
     NO = 0
     YES = 1
 
+class SingleValuedEnum(Enum):
+    VALUE = 1
+
 class ComparesEqualEnum(Enum):
     NO = 0
     YES = 1
@@ -70,13 +73,20 @@ class CustomNeEnum(Enum):
 
 static_assert(is_single_valued(Literal[NormalEnum.NO]))
 static_assert(is_single_valued(Literal[NormalEnum.YES]))
+static_assert(not is_single_valued(NormalEnum))
+
+static_assert(is_single_valued(Literal[SingleValuedEnum.VALUE]))
+static_assert(is_single_valued(SingleValuedEnum))
 
 static_assert(is_single_valued(Literal[ComparesEqualEnum.NO]))
 static_assert(is_single_valued(Literal[ComparesEqualEnum.YES]))
+static_assert(not is_single_valued(ComparesEqualEnum))
 
 static_assert(not is_single_valued(Literal[CustomEqEnum.NO]))
 static_assert(not is_single_valued(Literal[CustomEqEnum.YES]))
+static_assert(not is_single_valued(CustomEqEnum))
 
 static_assert(not is_single_valued(Literal[CustomNeEnum.NO]))
 static_assert(not is_single_valued(Literal[CustomNeEnum.YES]))
+static_assert(not is_single_valued(CustomNeEnum))
 ```

crates/ty_python_semantic/resources/mdtest/type_properties/is_singleton.md

@@ -13,11 +13,16 @@ class Answer(Enum):
     NO = 0
     YES = 1
 
+class Single(Enum):
+    VALUE = 1
+
 static_assert(is_singleton(None))
 static_assert(is_singleton(Literal[True]))
 static_assert(is_singleton(Literal[False]))
 static_assert(is_singleton(Literal[Answer.YES]))
 static_assert(is_singleton(Literal[Answer.NO]))
+static_assert(is_singleton(Literal[Single.VALUE]))
+static_assert(is_singleton(Single))
 
 static_assert(is_singleton(type[bool]))
 

crates/ty_python_semantic/resources/mdtest/type_properties/is_subtype_of.md

@@ -96,6 +96,9 @@ class Answer(Enum):
     NO = 0
     YES = 1
 
+class Single(Enum):
+    VALUE = 1
+
 # Boolean literals
 static_assert(is_subtype_of(Literal[True], bool))
 static_assert(is_subtype_of(Literal[True], int))
@@ -125,11 +128,12 @@ static_assert(is_subtype_of(Literal[b"foo"], object))
 static_assert(is_subtype_of(Literal[Answer.YES], Literal[Answer.YES]))
 static_assert(is_subtype_of(Literal[Answer.YES], Answer))
 static_assert(is_subtype_of(Literal[Answer.YES, Answer.NO], Answer))
-# TODO: this should not be an error
-# error: [static-assert-error]
 static_assert(is_subtype_of(Answer, Literal[Answer.YES, Answer.NO]))
 
 static_assert(not is_subtype_of(Literal[Answer.YES], Literal[Answer.NO]))
+
+static_assert(is_subtype_of(Literal[Single.VALUE], Single))
+static_assert(is_subtype_of(Single, Literal[Single.VALUE]))
 ```
 
 ## Heterogeneous tuple types

crates/ty_python_semantic/resources/mdtest/union_types.md

@@ -114,6 +114,33 @@ def _(
     reveal_type(u5)  # revealed: bool | Literal[17]
 ```
 
+## Enum literals
+
+```py
+from enum import Enum
+from typing import Literal
+
+class Color(Enum):
+    RED = "red"
+    GREEN = "green"
+    BLUE = "blue"
+
+def _(
+    u1: Literal[Color.RED, Color.GREEN],
+    u2: Color | Literal[Color.RED],
+    u3: Literal[Color.RED] | Color,
+    u4: Literal[Color.RED] | Literal[Color.RED, Color.GREEN],
+    u5: Literal[Color.RED, Color.GREEN, Color.BLUE],
+    u6: Literal[Color.RED] | Literal[Color.GREEN] | Literal[Color.BLUE],
+) -> None:
+    reveal_type(u1)  # revealed: Literal[Color.RED, Color.GREEN]
+    reveal_type(u2)  # revealed: Color
+    reveal_type(u3)  # revealed: Color
+    reveal_type(u4)  # revealed: Literal[Color.RED, Color.GREEN]
+    reveal_type(u5)  # revealed: Color
+    reveal_type(u6)  # revealed: Color
+```
+
 ## Do not erase `Unknown`
 
 ```py