mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-26 18:06:43 +00:00
dc66019fbc
## 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
2.8 KiB
2.8 KiB
Narrowing for is conditionals
is None
def _(flag: bool):
x = None if flag else 1
if x is None:
reveal_type(x) # revealed: None
else:
reveal_type(x) # revealed: Literal[1]
reveal_type(x) # revealed: None | Literal[1]
is for other types
def _(flag: bool):
class A: ...
x = A()
y = x if flag else None
if y is x:
reveal_type(y) # revealed: A
else:
reveal_type(y) # revealed: A | None
reveal_type(y) # revealed: A | None
is in chained comparisons
def _(x_flag: bool, y_flag: bool):
x = True if x_flag else False
y = True if y_flag else False
reveal_type(x) # revealed: bool
reveal_type(y) # revealed: bool
if y is x is False: # Interpreted as `(y is x) and (x is False)`
reveal_type(x) # revealed: Literal[False]
reveal_type(y) # revealed: bool
else:
# The negation of the clause above is (y is not x) or (x is not False)
# So we can't narrow the type of x or y here, because each arm of the `or` could be true
reveal_type(x) # revealed: bool
reveal_type(y) # revealed: bool
is in elif clause
def _(flag1: bool, flag2: bool):
x = None if flag1 else (1 if flag2 else True)
reveal_type(x) # revealed: None | Literal[1, True]
if x is None:
reveal_type(x) # revealed: None
elif x is True:
reveal_type(x) # revealed: Literal[True]
else:
reveal_type(x) # revealed: Literal[1]
is for enums
from enum import Enum
class Answer(Enum):
NO = 0
YES = 1
def _(answer: Answer):
if answer is 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 is Single.VALUE:
reveal_type(x) # revealed: Single
else:
reveal_type(x) # revealed: int
is for EllipsisType (Python 3.10+)
[environment]
python-version = "3.10"
from types import EllipsisType
def _(x: int | EllipsisType):
if x is ...:
reveal_type(x) # revealed: EllipsisType
else:
reveal_type(x) # revealed: int
is for EllipsisType (Python 3.9 and below)
[environment]
python-version = "3.9"
def _(flag: bool):
x = ... if flag else 42
reveal_type(x) # revealed: ellipsis | Literal[42]
if x is ...:
reveal_type(x) # revealed: ellipsis
else:
reveal_type(x) # revealed: Literal[42]
Assignment expressions
from typing import Literal
def f() -> Literal[1, 2] | None: ...
if (x := f()) is None:
reveal_type(x) # revealed: None
else:
reveal_type(x) # revealed: Literal[1, 2]