mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-27 18:36:35 +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
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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+)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue