[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
This commit is contained in:
David Peter 2025-07-21 19:37:55 +02:00 committed by GitHub
parent 926e83323a
commit dc66019fbc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 750 additions and 102 deletions

View file

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