mirror of
https://github.com/astral-sh/ruff.git
synced 2025-09-29 21:34:57 +00:00
[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:
parent
926e83323a
commit
dc66019fbc
19 changed files with 750 additions and 102 deletions
|
@ -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
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue