diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index c3b1e55c53e852..1458f63b8d08b1 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -1118,6 +1118,142 @@ 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 +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 `__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 26fc5620560de0..609d1803eae0f4 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 `Any` for the member's value type. + let custom_new_value_ty = if has_custom_enum_new(db, class) { + Some(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,53 @@ 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 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) { + 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 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; + } + + 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(_)) + }) +}