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

sharkdp · 2025-07-16T09:56:13Z

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.

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

github-actions · 2025-07-16T09:59:56Z

`mypy_primer` results

Changes were detected when running on open source projects

pwndbg (https://github.com/pwndbg/pwndbg)
- error[invalid-return-type] pwndbg/dbg/gdb/__init__.py:1660:10: Function can implicitly return `None`, which is not assignable to return type `((...) -> T, /) -> (...) -> T`
- Found 2270 diagnostics
+ Found 2269 diagnostics

urllib3 (https://github.com/urllib3/urllib3)
- error[invalid-return-type] src/urllib3/_collections.py:386:20: Return type does not match returned value: expected `list[str] | _DT`, found `(_Sentinel & ~Literal[_Sentinel.not_passed]) | (_DT & ~Literal[_Sentinel.not_passed])`
- error[invalid-argument-type] src/urllib3/util/connection.py:70:33: Argument to bound method `settimeout` is incorrect: Expected `int | float | None`, found `Unknown | _TYPE_DEFAULT`
- Found 402 diagnostics
+ Found 400 diagnostics

pytest (https://github.com/pytest-dev/pytest)
- error[invalid-argument-type] src/_pytest/fixtures.py:157:22: Argument to function `assert_never` is incorrect: Expected `Never`, found `Scope & ~Literal[Scope.Function] & ~Literal[Scope.Class] & ~Literal[Scope.Module] & ~Literal[Scope.Package] & ~Literal[Scope.Session]`
- error[invalid-argument-type] src/_pytest/fixtures.py:218:22: Argument to function `assert_never` is incorrect: Expected `Never`, found `Scope & ~Literal[Scope.Function] & ~Literal[Scope.Session] & ~Literal[Scope.Package] & ~Literal[Scope.Module] & ~Literal[Scope.Class]`
- error[invalid-argument-type] src/_pytest/pathlib.py:585:22: Argument to function `assert_never` is incorrect: Expected `Never`, found `ImportMode & ~Literal[ImportMode.importlib] & ~Literal[ImportMode.append] & ~Literal[ImportMode.prepend]`
- error[invalid-argument-type] testing/test_cacheprovider.py:1303:22: Argument to function `assert_never` is incorrect: Expected `Never`, found `Action & ~Literal[Action.MKDIR] & ~Literal[Action.SET]`
- error[invalid-argument-type] testing/test_cacheprovider.py:1315:22: Argument to function `assert_never` is incorrect: Expected `Never`, found `Action & ~Literal[Action.MKDIR] & ~Literal[Action.SET]`
- Found 517 diagnostics
+ Found 512 diagnostics

static-frame (https://github.com/static-frame/static-frame)
- warning[possibly-unresolved-reference] static_frame/core/join.py:328:15: Name `final_index` used when possibly not defined
- Found 1792 diagnostics
+ Found 1791 diagnostics

meson (https://github.com/mesonbuild/meson)
- warning[possibly-unresolved-reference] mesonbuild/mlog.py:287:16: Name `label` used when possibly not defined
- Found 912 diagnostics
+ Found 911 diagnostics

No memory usage changes detected ✅

github-actions · 2025-07-16T10:04:35Z

`ecosystem-analyzer` results

Lint rule	Added	Removed	Changed
`invalid-argument-type`	0	6	0
`invalid-return-type`	0	2	0
`possibly-unresolved-reference`	0	2	0
Total	0	10	0

Full report with detailed diff

codspeed-hq · 2025-07-16T10:06:19Z

CodSpeed Instrumentation Performance Report

Merging #19382 will not alter performance

_{Comparing david/enum-expansion (1252ef8) with main (5cace28)}

Summary

✅ 41 untouched benchmarks

carljm

Only got to reviewing the tests so far, but the tests look great!

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

crates/ty_python_semantic/src/types.rs

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

crates/ty_python_semantic/src/types/builder.rs

AlexWaygood · 2025-07-21T13:23:40Z

crates/ty_python_semantic/src/types/builder.rs

+                    // Note: we manually construct a `UnionType` here instead of going through
+                    // `UnionBuilder` because we would simplify the union to just the enum instance
+                    // and end up in this branch again.
+                    let db = self.db;
+                    self.add_positive(Type::Union(UnionType::new(
+                        db,
+                        enum_member_literals(db, instance.class.class_literal(db).0, None)
+                            .collect::<Box<[_]>>(),
+                    )))


oh, that's clever!

Initially, I duplicated what the Type::Union branch does here, but wanted to get rid of that. Fortunately, the wrapping/unwrapping in UnionType doesn't appear to hurt, performance-wise.

crates/ty_python_semantic/src/types/builder.rs

crates/ty_python_semantic/src/types/call/arguments.rs

…ability * origin/main: [ty] Expansion of enums into unions of literals (#19382) [ty] Avoid rechecking the entire project when changing the opened files (#19463) [ty] Add warning for unknown `TY_MEMORY_REPORT` value (#19465)

* main: (76 commits) Move fix suggestion to subdiagnostic (#19464) [ty] Implement non-stdlib stub mapping for classes and functions (#19471) [ty] Disallow illegal uses of `ClassVar` (#19483) [ty] Disallow `Final` in function parameter/return-type annotations (#19480) [ty] Extend `Final` test suite (#19476) [ty] Minor change to diagnostic message for invalid Literal uses (#19482) [ty] Detect illegal non-enum attribute accesses in Literal annotation (#19477) [ty] Reduce size of `TypeInference` (#19435) Run MD tests for Markdown-only changes (#19479) Revert "[ty] Detect illegal non-enum attribute accesses in Literal annotation" [ty] Detect illegal non-enum attribute accesses in Literal annotation [ty] Added semantic token support for more identifiers (#19473) [ty] Make tuple subclass constructors sound (#19469) [ty] Pass down specialization to generic dataclass bases (#19472) [ty] Garbage-collect reachability constraints (#19414) [ty] Implicit instance attributes declared `Final` (#19462) [ty] Expansion of enums into unions of literals (#19382) [ty] Avoid rechecking the entire project when changing the opened files (#19463) [ty] Add warning for unknown `TY_MEMORY_REPORT` value (#19465) [ty] Sync vendored typeshed stubs (#19461) ...

sharkdp added ty Multi-file analysis & type inference ecosystem-analyzer labels Jul 16, 2025

sharkdp force-pushed the david/enum-expansion branch from 07faf92 to 7351142 Compare July 17, 2025 08:54

sharkdp added ecosystem-analyzer and removed ecosystem-analyzer labels Jul 17, 2025

sharkdp added 10 commits July 18, 2025 12:17

[ty] Expansion of enums into unions of literals

ea62f9a

Enum expansion in overload resolution

e9a886b

New tests

b891bf1

Some bug fixes

d72a6ba

Fix single-enum member behavior

1c3a47f

More fixes for single-member enums

3490022

Fix equality narrowing

e1ea489

More exhaustiveness tests

b20a243

Fix single-valuedness

c11770e

FxOrderSet

d8ada59

sharkdp force-pushed the david/enum-expansion branch from 9f8c78d to d8ada59 Compare July 18, 2025 11:46

Cleanup

4a6f195

sharkdp added ecosystem-analyzer and removed ecosystem-analyzer labels Jul 18, 2025

sharkdp marked this pull request as ready for review July 18, 2025 12:02

sharkdp requested a review from carljm as a code owner July 18, 2025 12:02

sharkdp requested review from AlexWaygood and dcreager as code owners July 18, 2025 12:02

carljm reviewed Jul 19, 2025

View reviewed changes

crates/ty_python_semantic/resources/mdtest/narrow/conditionals/eq.md Show resolved Hide resolved

AlexWaygood approved these changes Jul 21, 2025

View reviewed changes

sharkdp added 7 commits July 21, 2025 15:56

Minor simplification

952f04d

New is_equivalent_to tests

e8f4b1b

UnionElement::to_type_element

804bdce

Functional style

72648bb

Change return type of enum_member_literals

312f83f

Add test for non-expansion of Enum subclasses without members

74a2798

Clippy

1252ef8

sharkdp merged commit dc66019 into main Jul 21, 2025
37 checks passed

sharkdp deleted the david/enum-expansion branch July 21, 2025 17:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

sharkdp commented Jul 16, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 16, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 16, 2025 •

edited

Loading

Uh oh!

codspeed-hq bot commented Jul 16, 2025 •

edited

Loading

Uh oh!

carljm left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AlexWaygood Jul 21, 2025

Uh oh!

sharkdp Jul 21, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

sharkdp commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Performance

Test Plan

Uh oh!

github-actions bot commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

mypy_primer results

Uh oh!

github-actions bot commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ecosystem-analyzer results

Uh oh!

codspeed-hq bot commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CodSpeed Instrumentation Performance Report

Merging #19382 will not alter performance

Summary

Uh oh!

carljm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AlexWaygood Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

sharkdp Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sharkdp commented Jul 16, 2025 •

edited

Loading

github-actions bot commented Jul 16, 2025 •

edited

Loading

`mypy_primer` results

github-actions bot commented Jul 16, 2025 •

edited

Loading

`ecosystem-analyzer` results

codspeed-hq bot commented Jul 16, 2025 •

edited

Loading