diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 43eef415753de..59a04b1d96dc7 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -694,6 +694,51 @@ class Child(Base): GITHUB = "github" # error: [invalid-assignment] ``` +### Inherited `__new_member__` + +An inherited `__new_member__` takes precedence over the default enum constructor and can replace the +member value: + +```py +from enum import Enum + +class Base(Enum): + def __new_member__(cls: type["Base"], value: int) -> "Base": + obj = object.__new__(cls) + obj._value_ = str(value) + return obj + +class Child(Base): + VALUE = 1 + +reveal_type(Child.VALUE.value) # revealed: Any +``` + +`EnumType` saves an enum's user-defined `__new__` as that class's `__new_member__`. The immediate +parent's `__new__` therefore takes precedence over an explicit `__new_member__` in a grandparent: + +```py +from enum import Enum + +class Grandparent(Enum): + def __new_member__(cls: type["Grandparent"], value: str) -> "Grandparent": + obj = object.__new__(cls) + obj._value_ = value + return obj + +class Parent(Grandparent): + def __new__(cls, value: int) -> "Parent": + obj = object.__new__(cls) + obj._value_ = value + return obj + +class Child(Parent): + VALID = 1 + INVALID = "not an int" # error: [invalid-assignment] + +reveal_type(Child.VALID.value) # revealed: Any +``` + ### Data-type mixin `__new__` A user-defined `__new__` on a data-type mixin constructs the scalar payload and can transform the diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index c7b3e4a8b4dc8..9eb007d8d966b 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -872,8 +872,10 @@ pub(crate) fn enum_metadata<'db>( let init = resolve_enum_method(user_defined_init, || { inherited_known_enum_method(db, class, "__init__") }); + // CPython checks `__new_member__` and then `__new__` on each enum base before continuing + // through the MRO or falling back to the data-type constructor. let user_defined_new = custom_enum_method(db, scope_id, "__new__") - .or_else(|| inherited_user_defined_enum_method(db, class, "__new__")) + .or_else(|| inherited_user_defined_enum_new(db, class)) .or_else(|| inherited_user_defined_mixin_new(db, class)); let new = resolve_enum_method(user_defined_new, || { inherited_known_enum_method(db, class, "__new__") @@ -1216,6 +1218,20 @@ fn inherited_user_defined_enum_method<'db>( .find_map(|base| custom_enum_method(db, base.body_scope(db), name)) } +/// Looks up the first user-defined enum member constructor in the MRO. +fn inherited_user_defined_enum_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| { + let scope = base.body_scope(db); + custom_enum_method(db, scope, "__new_member__") + .or_else(|| custom_enum_method(db, scope, "__new__")) + }) +} + /// Looks up a user-defined `__new__` on a data-type mixin anywhere in the MRO, including through an /// enum base. ///