From 3270cd838012dc10af66447c91fa12cc21c8760c Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 17 May 2026 19:18:08 +0200 Subject: [PATCH 1/2] [ty] Ignore _generate_next_value_ with custom construction hooks --- .../resources/mdtest/enums.md | 22 +++++++++++++ crates/ty_python_semantic/src/types/enums.rs | 33 ++++++++++++++----- 2 files changed, 46 insertions(+), 9 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 3df4884824c7e..b62202f35805b 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -1169,6 +1169,7 @@ precedence: ```py from enum import Enum, EnumMeta, IntEnum, auto +from ty_extensions import enum_members from typing import Literal class WithNewAndGenerateNextValue(Enum): @@ -1190,6 +1191,27 @@ reveal_type(WithNewAndGenerateNextValue.A.value) # revealed: Any def _instance_new(a: WithNewAndGenerateNextValue): reveal_type(a.value) # revealed: Any +class WithNewAndLiteralGenerateNextValue(Enum): + @staticmethod + def _generate_next_value_(name, start, count, last_values) -> Literal["x"]: + return "x" + + def __new__(cls, value: str) -> "WithNewAndLiteralGenerateNextValue": + obj = object.__new__(cls) + obj._value_ = object() + return obj + + A = auto() + B = auto() + +# `__new__` can rewrite duplicate generated values to distinct values, so `B` is not an alias of `A`. +# revealed: tuple[Literal["A"], Literal["B"]] +reveal_type(enum_members(WithNewAndLiteralGenerateNextValue)) +reveal_type(WithNewAndLiteralGenerateNextValue.A) # revealed: Literal[WithNewAndLiteralGenerateNextValue.A] +reveal_type(WithNewAndLiteralGenerateNextValue.B) # revealed: Literal[WithNewAndLiteralGenerateNextValue.B] +reveal_type(WithNewAndLiteralGenerateNextValue.A.value) # revealed: Any +reveal_type(WithNewAndLiteralGenerateNextValue.B.value) # revealed: Any + class WithInitAndGenerateNextValue(Enum): @staticmethod def _generate_next_value_(name, start, count, last_values) -> str: diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 7c6516443012e..0edf303765a89 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -218,16 +218,27 @@ fn try_register_alias<'db>( /// with a custom `_generate_next_value_`, aliasing is based on the generated /// value instead of the pre-generator placeholder used while collecting /// members. +/// +/// Returns `None` for `auto()` members when construction hooks can rewrite +/// `_value_`, because neither the generated value nor the placeholder is +/// reliable alias evidence in that case. fn alias_detection_value<'db>( db: &'db dyn Db, value_ty: Type<'db>, is_auto: bool, generate_next_value_function: Option>, -) -> Type<'db> { - if is_auto && let Some(func_ty) = generate_next_value_function { - func_ty.signature(db).overload_return_type_or_unknown(db) + init_function: Option>, + new_function: Option>, + custom_enum_metaclass_new: bool, +) -> Option> { + if !is_auto { + Some(value_ty) + } else if init_function.is_some() || new_function.is_some() || custom_enum_metaclass_new { + None + } else if let Some(func_ty) = generate_next_value_function { + Some(func_ty.signature(db).overload_return_type_or_unknown(db)) } else { - value_ty + Some(value_ty) } } @@ -301,6 +312,9 @@ pub(crate) fn enum_metadata<'db>( let mut prev_value_was_non_literal_int = false; let mut prev_bool_literal = None; let ignored_names = enum_ignored_names(db, scope_id); + let init_function = custom_init(db, scope_id).or_else(|| inherited_init(db, class)); + let new_function = custom_new(db, scope_id).or_else(|| inherited_new(db, class)); + let custom_enum_metaclass_new = custom_enum_metaclass_new(db, class); let generate_next_value_function = custom_generate_next_value(db, scope_id) .or_else(|| inherited_generate_next_value(db, class)); @@ -444,8 +458,13 @@ pub(crate) fn enum_metadata<'db>( value_ty, auto_members.contains(name), generate_next_value_function, + init_function, + new_function, + custom_enum_metaclass_new, ); - if try_register_alias(alias_value_ty, name, &mut enum_values, &mut aliases) { + if let Some(alias_value_ty) = alias_value_ty + && try_register_alias(alias_value_ty, name, &mut enum_values, &mut aliases) + { return None; } @@ -488,10 +507,6 @@ pub(crate) fn enum_metadata<'db>( return None; } - // Look up custom construction hooks, falling back to parent enum classes. - let init_function = custom_init(db, scope_id).or_else(|| inherited_init(db, class)); - let new_function = custom_new(db, scope_id).or_else(|| inherited_new(db, class)); - let custom_enum_metaclass_new = custom_enum_metaclass_new(db, class); let custom_value_annotation = custom_value_annotation(db, scope_id); let value_annotation = custom_value_annotation.or_else(|| { if custom_enum_metaclass_new { From 4d7e503a2faca0255e7746bc31839c6b1293b2c3 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 17 May 2026 19:27:48 +0200 Subject: [PATCH 2/2] Only apply to custom news --- .../resources/mdtest/enums.md | 26 ++++++++++++++++ crates/ty_python_semantic/src/types/enums.rs | 30 +++++++++++++------ 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index b62202f35805b..2fc212118e62f 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -1227,6 +1227,20 @@ reveal_type(WithInitAndGenerateNextValue.A.value) # revealed: Any def _instance_init(a: WithInitAndGenerateNextValue): reveal_type(a.value) # revealed: Any +class WithInitAndLiteralGenerateNextValue(Enum): + @staticmethod + def _generate_next_value_(name, start, count, last_values) -> Literal["x"]: + return "x" + + def __init__(self, value: str) -> None: ... + + A = auto() + B = auto() + +# `__init__` runs after duplicate generated values are resolved to aliases. +# revealed: tuple[Literal["A"]] +reveal_type(enum_members(WithInitAndLiteralGenerateNextValue)) + class ChoicesType(EnumMeta): def __new__(metacls, classname, bases, classdict, **kwds): ... @@ -1244,6 +1258,18 @@ reveal_type(MyModelChoices.A.value) # revealed: Any def _instance_metaclass(a: MyModelChoices): reveal_type(a.value) # revealed: Any + +class IntEnumDuplicateAutoAliases(IntEnum): + @staticmethod + def _generate_next_value_(name, start, count, last_values) -> Literal[42]: + return 42 + + A = auto() + B = auto() + +# The stdlib `IntEnum.__new__` preserves duplicate generated values as aliases. +# revealed: tuple[Literal["A"]] +reveal_type(enum_members(IntEnumDuplicateAutoAliases)) ``` For non-`auto()` members in a mixed enum, `_generate_next_value_` does not apply at all, and the diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 0edf303765a89..bf9d88a2b81a6 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -219,21 +219,20 @@ fn try_register_alias<'db>( /// value instead of the pre-generator placeholder used while collecting /// members. /// -/// Returns `None` for `auto()` members when construction hooks can rewrite -/// `_value_`, because neither the generated value nor the placeholder is -/// reliable alias evidence in that case. +/// Returns `None` for `auto()` members when `__new__` or a custom metaclass can +/// rewrite `_value_` before alias registration, because neither the generated +/// value nor the placeholder is reliable alias evidence in that case. fn alias_detection_value<'db>( db: &'db dyn Db, value_ty: Type<'db>, is_auto: bool, generate_next_value_function: Option>, - init_function: Option>, - new_function: Option>, + user_defined_new_function: Option>, custom_enum_metaclass_new: bool, ) -> Option> { if !is_auto { Some(value_ty) - } else if init_function.is_some() || new_function.is_some() || custom_enum_metaclass_new { + } else if user_defined_new_function.is_some() || custom_enum_metaclass_new { None } else if let Some(func_ty) = generate_next_value_function { Some(func_ty.signature(db).overload_return_type_or_unknown(db)) @@ -312,8 +311,12 @@ pub(crate) fn enum_metadata<'db>( let mut prev_value_was_non_literal_int = false; let mut prev_bool_literal = None; let ignored_names = enum_ignored_names(db, scope_id); + + // Look up custom construction hooks, falling back to parent enum classes. let init_function = custom_init(db, scope_id).or_else(|| inherited_init(db, class)); - let new_function = custom_new(db, scope_id).or_else(|| inherited_new(db, class)); + let user_defined_new_function = + custom_new(db, scope_id).or_else(|| inherited_user_defined_new(db, class)); + let new_function = user_defined_new_function.or_else(|| inherited_new(db, class)); let custom_enum_metaclass_new = custom_enum_metaclass_new(db, class); let generate_next_value_function = custom_generate_next_value(db, scope_id) .or_else(|| inherited_generate_next_value(db, class)); @@ -458,8 +461,7 @@ pub(crate) fn enum_metadata<'db>( value_ty, auto_members.contains(name), generate_next_value_function, - init_function, - new_function, + user_defined_new_function, custom_enum_metaclass_new, ); if let Some(alias_value_ty) = alias_value_ty @@ -614,6 +616,16 @@ fn inherited_new<'db>( iter_parent_enum_classes(db, class).find_map(|base| custom_new(db, base.body_scope(db))) } +/// Looks up an inherited `__new__` from user-defined parent enum classes in the MRO. +fn inherited_user_defined_new<'db>( + db: &'db dyn Db, + class: StaticClassLiteral<'db>, +) -> Option> { + iter_parent_enum_classes(db, class) + .filter(|base| base.known(db).is_none()) + .find_map(|base| custom_new(db, base.body_scope(db))) +} + /// Looks up an inherited `_generate_next_value_` from parent enum classes in the MRO. fn inherited_generate_next_value<'db>( db: &'db dyn Db,