From 0bc125d01dc732453d3eff260245298f4f892135 Mon Sep 17 00:00:00 2001 From: Ray Zeller Date: Fri, 20 Feb 2026 09:39:58 -0700 Subject: [PATCH 1/2] [ty] Support custom `__new__` in enums for value type inference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When an enum defines a custom `__new__`, the raw assignment type doesn't represent the member's value — `__new__` unpacks arguments and explicitly sets `_value_`. Fall back to the `_value_` annotation type if declared, or `Any` otherwise. Closes: astral-sh/ty#876 (partial — "Custom `__new__` or `__init__` methods in enums") Co-Authored-By: Claude Opus 4.6 --- .../resources/mdtest/enums.md | 158 ++++++++++++++++++ crates/ty_python_semantic/src/types/enums.rs | 97 ++++++++++- 2 files changed, 254 insertions(+), 1 deletion(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index c3b1e55c53e85..1328b4004bd50 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -1118,6 +1118,164 @@ class MyEnum[T](MyEnumBase): A = 1 ``` +## Custom `__new__` + +When an enum defines a custom `__new__` method, the raw assignment values are unpacked as arguments +to `__new__`, and `_value_` is explicitly set inside the method body. The member's `.value` type is +determined by the `_value_` annotation (if declared) or falls back to `Any`. + +### Custom `__new__` with tuple values + +```py +from enum import Enum + +class Planet(Enum): + def __new__(cls, mass: float, radius: float) -> "Planet": + obj = object.__new__(cls) + obj._value_ = mass + return obj + MERCURY = (3.303e23, 2.4397e6) + VENUS = (4.869e24, 6.0518e6) + EARTH = (5.976e24, 6.37814e6) + +# Without a `_value_` annotation, values fall back to `Any` +reveal_type(Planet.MERCURY.value) # revealed: Any +reveal_type(Planet.VENUS.value) # revealed: Any +reveal_type(Planet.EARTH.value) # revealed: Any + +# `.name` still works correctly +reveal_type(Planet.MERCURY.name) # revealed: Literal["MERCURY"] +``` + +### Custom `__new__` with `_value_` annotation + +```py +from enum import Enum + +class Planet(Enum): + _value_: float + + def __new__(cls, mass: float, radius: float) -> "Planet": + obj = object.__new__(cls) + obj._value_ = mass + return obj + MERCURY = (3.303e23, 2.4397e6) + VENUS = (4.869e24, 6.0518e6) + EARTH = (5.976e24, 6.37814e6) + +# With a `_value_: float` annotation, values are inferred as `int | float` +# (`float` widens to `int | float` per PEP 484 numeric tower rules) +reveal_type(Planet.MERCURY.value) # revealed: int | float +reveal_type(Planet.VENUS.value) # revealed: int | float +``` + +### Custom `__init__` without `__new__` + +Custom `__init__` does **not** change `_value_` — the raw assignment is the value. + +```py +from enum import Enum + +class Mood(Enum): + def __init__(self, value: int): + self._mood_level = value + HAPPY = 1 + SAD = 2 + NEUTRAL = 3 + +reveal_type(Mood.HAPPY.value) # revealed: Literal[1] +reveal_type(Mood.SAD.value) # revealed: Literal[2] +``` + +### Inherited custom `__new__` from a base class + +```py +from enum import Enum + +class AutoNameEnum(Enum): + def __new__(cls, display_name: str) -> "AutoNameEnum": + obj = object.__new__(cls) + obj._value_ = display_name.lower() + return obj + +class Color(AutoNameEnum): + RED = "Red" + GREEN = "Green" + BLUE = "Blue" + +reveal_type(Color.RED.value) # revealed: Any +reveal_type(Color.GREEN.value) # revealed: Any +``` + +### Standard enums are not affected + +```py +from enum import Enum + +class Color(Enum): + RED = 1 + GREEN = 2 + BLUE = 3 + +reveal_type(Color.RED.value) # revealed: Literal[1] +reveal_type(Color.GREEN.value) # revealed: Literal[2] +reveal_type(Color.BLUE.value) # revealed: Literal[3] +``` + +### `IntEnum` not affected by `int.__new__` + +```py +from enum import IntEnum + +class Status(IntEnum): + ACTIVE = 1 + INACTIVE = 2 + +reveal_type(Status.ACTIVE.value) # revealed: Literal[1] +reveal_type(Status.INACTIVE.value) # revealed: Literal[2] +``` + +### Member detection still works with custom `__new__` + +```py +from enum import Enum +from typing_extensions import reveal_type + +class Planet(Enum): + def __new__(cls, mass: float, radius: float) -> "Planet": + obj = object.__new__(cls) + obj._value_ = mass + return obj + MERCURY = (3.303e23, 2.4397e6) + VENUS = (4.869e24, 6.0518e6) + +reveal_type(Planet.MERCURY) # revealed: Literal[Planet.MERCURY] +reveal_type(Planet.VENUS) # revealed: Literal[Planet.VENUS] +``` + +### Exhaustiveness with custom `__new__` + +```py +from enum import Enum +from typing_extensions import assert_never + +class Planet(Enum): + def __new__(cls, mass: float, radius: float) -> "Planet": + obj = object.__new__(cls) + obj._value_ = mass + return obj + MERCURY = (3.303e23, 2.4397e6) + VENUS = (4.869e24, 6.0518e6) + +def check(planet: Planet) -> str: + if planet is Planet.MERCURY: + return "Mercury" + elif planet is Planet.VENUS: + return "Venus" + else: + assert_never(planet) +``` + ## References - Typing spec: diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 26fc5620560de..1faf19fcdab76 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -79,6 +79,15 @@ pub(crate) fn enum_metadata<'db>( let use_def_map = use_def_map(db, scope_id); let table = place_table(db, scope_id); + // When an enum has a custom `__new__`, the raw assignment type doesn't represent the + // member's value — `__new__` unpacks the arguments and explicitly sets `_value_`. + // Fall back to the `_value_` annotation type (if declared) or `Any`. + let custom_new_value_ty = if has_custom_enum_new(db, class) { + Some(enum_value_annotation_type(db, class).unwrap_or(Type::any())) + } else { + None + }; + let mut enum_values: FxHashMap, Name> = FxHashMap::default(); let mut auto_counter = 0; @@ -276,7 +285,8 @@ pub(crate) fn enum_metadata<'db>( } } - Some((name.clone(), value_ty)) + let final_value_ty = custom_new_value_ty.unwrap_or(value_ty); + Some((name.clone(), final_value_ty)) }) .collect::>(); @@ -348,3 +358,88 @@ pub(crate) fn try_unwrap_nonmember_value<'db>(db: &'db dyn Db, ty: Type<'db>) -> _ => None, } } + +/// Returns `true` if the enum class (or a class in its MRO) defines a custom `__new__` method. +/// +/// When an enum has a custom `__new__`, the assigned tuple values are unpacked as arguments to +/// `__new__`, and `_value_` is explicitly set inside the method body. This means we can't infer +/// the member's value type from the raw assignment — we need to fall back to the `_value_` +/// annotation (if any) or `Any`. +/// +/// We skip classes that are expected to have standard `__new__` implementations: +/// `object`, `Enum`, `StrEnum`, `int`, `str`, `float`, `complex`, `bytes`, and `bool`. +fn has_custom_enum_new<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> bool { + // Check the enum class itself + if has_own_dunder_new(db, class) { + return true; + } + + // Walk the MRO (skipping the class itself) looking for a custom `__new__` + class + .iter_mro(db, None) + .skip(1) + .filter_map(ClassBase::into_class) + .any(|mro_class| { + let Some(static_class) = mro_class.class_literal(db).as_static() else { + return false; + }; + + // Skip known classes with standard `__new__` implementations + if matches!( + static_class.known(db), + Some( + KnownClass::Object + | KnownClass::Enum + | KnownClass::StrEnum + | KnownClass::Int + | KnownClass::Str + | KnownClass::Float + | KnownClass::Complex + | KnownClass::Bytes + | KnownClass::Bool + ) + ) { + return false; + } + + // Skip classes defined in stub files (e.g. `IntEnum.__new__`, `IntFlag.__new__` + // from typeshed). These `__new__` definitions exist for typing purposes and + // don't represent custom value transformations. + if static_class.body_scope(db).file(db).is_stub(db) { + return false; + } + + has_own_dunder_new(db, static_class) + }) +} + +/// Returns `true` if the class defines `__new__` directly in its own body scope. +fn has_own_dunder_new<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> bool { + let scope = class.body_scope(db); + let table = place_table(db, scope); + table.symbol_id("__new__").is_some_and(|symbol_id| { + let bindings = use_def_map(db, scope).reachable_symbol_bindings(symbol_id); + matches!(place_from_bindings(db, bindings).place, Place::Defined(_)) + }) +} + +/// When an enum class has a custom `__new__`, look up the `_value_` annotation +/// on the class itself to determine the value type. Returns `None` if no annotation +/// is found on the enum class's own body. +fn enum_value_annotation_type<'db>( + db: &'db dyn Db, + class: StaticClassLiteral<'db>, +) -> Option> { + let scope_id = class.body_scope(db); + let use_def = use_def_map(db, scope_id); + let table = place_table(db, scope_id); + + let symbol_id = table.symbol_id("_value_")?; + let declarations = use_def.end_of_scope_symbol_declarations(symbol_id); + let declared = place_from_declarations(db, declarations).ignore_conflicting_declarations(); + + match declared.place { + Place::Defined(DefinedPlace { ty, .. }) if !ty.is_dynamic() => Some(ty), + _ => None, + } +} From ab0238a1c718afa49a82e47f05f9ae2b76d48d6e Mon Sep 17 00:00:00 2001 From: Ray Zeller Date: Fri, 20 Feb 2026 13:49:17 -0700 Subject: [PATCH 2/2] [ty] Strip `_value_` annotation lookup and simplify `has_custom_enum_new` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the `_value_` annotation lookup logic to avoid overlap with #22228, which handles `_value_` semantics more broadly. Custom `__new__` enums now always fall back to `Any`. Also remove the redundant known-class skip list from `has_custom_enum_new` — the vendored path check already covers all stdlib classes. Co-Authored-By: Claude Opus 4.6 --- .../resources/mdtest/enums.md | 26 +------- crates/ty_python_semantic/src/types/enums.rs | 65 +++++-------------- 2 files changed, 17 insertions(+), 74 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 1328b4004bd50..1458f63b8d08b 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -1121,8 +1121,8 @@ class MyEnum[T](MyEnumBase): ## Custom `__new__` When an enum defines a custom `__new__` method, the raw assignment values are unpacked as arguments -to `__new__`, and `_value_` is explicitly set inside the method body. The member's `.value` type is -determined by the `_value_` annotation (if declared) or falls back to `Any`. +to `__new__`, and `_value_` is explicitly set inside the method body. The member's `.value` type +falls back to `Any`. ### Custom `__new__` with tuple values @@ -1147,28 +1147,6 @@ reveal_type(Planet.EARTH.value) # revealed: Any reveal_type(Planet.MERCURY.name) # revealed: Literal["MERCURY"] ``` -### Custom `__new__` with `_value_` annotation - -```py -from enum import Enum - -class Planet(Enum): - _value_: float - - def __new__(cls, mass: float, radius: float) -> "Planet": - obj = object.__new__(cls) - obj._value_ = mass - return obj - MERCURY = (3.303e23, 2.4397e6) - VENUS = (4.869e24, 6.0518e6) - EARTH = (5.976e24, 6.37814e6) - -# With a `_value_: float` annotation, values are inferred as `int | float` -# (`float` widens to `int | float` per PEP 484 numeric tower rules) -reveal_type(Planet.MERCURY.value) # revealed: int | float -reveal_type(Planet.VENUS.value) # revealed: int | float -``` - ### Custom `__init__` without `__new__` Custom `__init__` does **not** change `_value_` — the raw assignment is the value. diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 1faf19fcdab76..609d1803eae0f 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -81,9 +81,9 @@ pub(crate) fn enum_metadata<'db>( // When an enum has a custom `__new__`, the raw assignment type doesn't represent the // member's value — `__new__` unpacks the arguments and explicitly sets `_value_`. - // Fall back to the `_value_` annotation type (if declared) or `Any`. + // Fall back to `Any` for the member's value type. let custom_new_value_ty = if has_custom_enum_new(db, class) { - Some(enum_value_annotation_type(db, class).unwrap_or(Type::any())) + Some(Type::any()) } else { None }; @@ -363,11 +363,7 @@ pub(crate) fn try_unwrap_nonmember_value<'db>(db: &'db dyn Db, ty: Type<'db>) -> /// /// When an enum has a custom `__new__`, the assigned tuple values are unpacked as arguments to /// `__new__`, and `_value_` is explicitly set inside the method body. This means we can't infer -/// the member's value type from the raw assignment — we need to fall back to the `_value_` -/// annotation (if any) or `Any`. -/// -/// We skip classes that are expected to have standard `__new__` implementations: -/// `object`, `Enum`, `StrEnum`, `int`, `str`, `float`, `complex`, `bytes`, and `bool`. +/// the member's value type from the raw assignment — we fall back to `Any`. fn has_custom_enum_new<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> bool { // Check the enum class itself if has_own_dunder_new(db, class) { @@ -384,28 +380,18 @@ fn has_custom_enum_new<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> return false; }; - // Skip known classes with standard `__new__` implementations - if matches!( - static_class.known(db), - Some( - KnownClass::Object - | KnownClass::Enum - | KnownClass::StrEnum - | KnownClass::Int - | KnownClass::Str - | KnownClass::Float - | KnownClass::Complex - | KnownClass::Bytes - | KnownClass::Bool - ) - ) { - return false; - } - - // Skip classes defined in stub files (e.g. `IntEnum.__new__`, `IntFlag.__new__` - // from typeshed). These `__new__` definitions exist for typing purposes and - // don't represent custom value transformations. - if static_class.body_scope(db).file(db).is_stub(db) { + // Skip classes defined in vendored typeshed (e.g. `object.__new__`, + // `int.__new__`, `IntEnum.__new__`, `IntFlag.__new__`). These `__new__` + // definitions exist for typing purposes and don't represent custom value + // transformations. We specifically check for vendored paths rather than all + // stub files, because a user-provided `.pyi` stub for a library with a + // custom `__new__` should still be recognized. + if static_class + .body_scope(db) + .file(db) + .path(db) + .is_vendored_path() + { return false; } @@ -422,24 +408,3 @@ fn has_own_dunder_new<'db>(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> b matches!(place_from_bindings(db, bindings).place, Place::Defined(_)) }) } - -/// When an enum class has a custom `__new__`, look up the `_value_` annotation -/// on the class itself to determine the value type. Returns `None` if no annotation -/// is found on the enum class's own body. -fn enum_value_annotation_type<'db>( - db: &'db dyn Db, - class: StaticClassLiteral<'db>, -) -> Option> { - let scope_id = class.body_scope(db); - let use_def = use_def_map(db, scope_id); - let table = place_table(db, scope_id); - - let symbol_id = table.symbol_id("_value_")?; - let declarations = use_def.end_of_scope_symbol_declarations(symbol_id); - let declared = place_from_declarations(db, declarations).ignore_conflicting_declarations(); - - match declared.place { - Place::Defined(DefinedPlace { ty, .. }) if !ty.is_dynamic() => Some(ty), - _ => None, - } -}