Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 45 additions & 0 deletions crates/ty_python_semantic/resources/mdtest/enums.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 17 additions & 1 deletion crates/ty_python_semantic/src/types/enums.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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__")
Expand Down Expand Up @@ -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<EnumMethodBinding<'db>> {
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.
///
Expand Down
Loading