Validate against MRO
diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 3ca761f..8f7088c 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md
@@ -141,6 +141,26 @@ from enum import Enum class Planet(Enum): + _value_: int + + def __init__(self, value: int, mass: float, radius: float): + self._value_ = value + + MERCURY = (1, 3.303e23, 2.4397e6) + SATURN = "saturn" # error: [invalid-assignment] + +reveal_type(Planet.MERCURY.value) # revealed: int +reveal_type(Planet.MERCURY._value_) # revealed: int +``` + +### `_value_` annotation incompatible with `__init__` + +When `_value_` and `__init__` disagree, the assignment inside `__init__` is flagged: + +```py +from enum import Enum + +class Planet(Enum): _value_: str def __init__(self, value: int, mass: float, radius: float): @@ -176,8 +196,8 @@ ### Inherited `_value_` annotation -A `_value_` annotation on a parent enum is not inherited by subclasses for the purpose of member -value validation: +A `_value_` annotation on a parent enum is inherited by subclasses. Member values are validated +against the inherited annotation, and `.value` uses the declared type: ```py from enum import Enum @@ -187,9 +207,66 @@ class Child(Base): A = 1 - B = "not checked against int" + B = "not an int" # error: [invalid-assignment] -reveal_type(Child.A.value) # revealed: Literal[1] +reveal_type(Child.A.value) # revealed: int +``` + +This also works through multiple levels of inheritance, where `_value_` is declared on an +intermediate class: + +```py +from enum import Enum + +class Grandparent(Enum): + pass + +class Parent(Grandparent): + _value_: int + +class Child(Parent): + A = 1 + B = "not an int" # error: [invalid-assignment] + +reveal_type(Child.A.value) # revealed: int +``` + +### Inherited `__init__` + +A custom `__init__` on a parent enum is inherited by subclasses. Member values are validated against +the inherited `__init__` signature: + +```py +from enum import Enum + +class Base(Enum): + def __init__(self, a: int, b: str): + self._value_ = a + +class Child(Base): + A = (1, "foo") + B = "should be checked against __init__" # error: [invalid-assignment] + +reveal_type(Child.A.value) # revealed: Any +``` + +This also works through multiple levels of inheritance: + +```py +from enum import Enum + +class Grandparent(Enum): + def __init__(self, a: int, b: str): + self._value_ = a + +class Parent(Grandparent): + pass + +class Child(Parent): + A = (1, "foo") + B = "bad" # error: [invalid-assignment] + +reveal_type(Child.A.value) # revealed: Any ``` ### Non-member attributes with disallowed type @@ -450,7 +527,8 @@ reveal_type(SingleMember.SINGLE.value) # revealed: Literal["single"] ``` -Using `auto()` with `IntEnum` also works as expected: +Using `auto()` with `IntEnum` also works as expected. `IntEnum` declares `_value_: int` in typeshed, +so `.value` is typed as `int` rather than a precise literal: ```py from enum import IntEnum, auto @@ -459,8 +537,8 @@ YES = auto() NO = auto() -reveal_type(Answer.YES.value) # revealed: Literal[1] -reveal_type(Answer.NO.value) # revealed: Literal[2] +reveal_type(Answer.YES.value) # revealed: int +reveal_type(Answer.NO.value) # revealed: int ``` As does using `auto()` for other enums that use `int` as a mixin:
diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 2ecf7e1..0cb4eb7 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs
@@ -315,7 +315,8 @@ return None; } - // Look up an explicit `_value_` annotation, if present. + // Look up an explicit `_value_` annotation, if present. Falls back to + // checking parent enum classes in the MRO. let value_annotation = place_table(db, scope_id) .symbol_id("_value_") .and_then(|symbol_id| { @@ -323,9 +324,11 @@ place_from_declarations(db, declarations) .ignore_conflicting_declarations() .ignore_possibly_undefined() - }); + }) + .or_else(|| inherited_value_annotation(db, class)); - let init_function = custom_init(db, scope_id); + // Look up a custom `__init__`, falling back to parent enum classes. + let init_function = custom_init(db, scope_id).or_else(|| inherited_init(db, class)); Some(EnumMetadata { members, @@ -335,6 +338,56 @@ }) } +/// Iterates over parent enum classes in the MRO, skipping known classes +/// (like `Enum`, `StrEnum`, etc.) that we handle specially. +fn iter_parent_enum_classes<'db>( + db: &'db dyn Db, + class: StaticClassLiteral<'db>, +) -> impl Iterator<Item = StaticClassLiteral<'db>> + 'db { + class + .iter_mro(db, None) + .skip(1) + .filter_map(ClassBase::into_class) + .filter_map(move |class_type| { + let base = class_type.class_literal(db).as_static()?; + (base.known(db).is_none() && is_enum_class_by_inheritance(db, base)).then_some(base) + }) +} + +/// Looks up an inherited `_value_` annotation from parent enum classes in the MRO. +fn inherited_value_annotation<'db>( + db: &'db dyn Db, + class: StaticClassLiteral<'db>, +) -> Option<Type<'db>> { + for base_class in iter_parent_enum_classes(db, class) { + let scope_id = base_class.body_scope(db); + let use_def = use_def_map(db, scope_id); + if let Some(symbol_id) = place_table(db, scope_id).symbol_id("_value_") { + let declarations = use_def.end_of_scope_symbol_declarations(symbol_id); + if let Some(ty) = place_from_declarations(db, declarations) + .ignore_conflicting_declarations() + .ignore_possibly_undefined() + { + return Some(ty); + } + } + } + None +} + +/// Looks up an inherited `__init__` from parent enum classes in the MRO. +fn inherited_init<'db>( + db: &'db dyn Db, + class: StaticClassLiteral<'db>, +) -> Option<FunctionType<'db>> { + for base_class in iter_parent_enum_classes(db, class) { + if let Some(f) = custom_init(db, base_class.body_scope(db)) { + return Some(f); + } + } + None +} + /// Returns the custom `__init__` function type if one is defined on the enum. fn custom_init<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option<FunctionType<'db>> { let init_symbol_id = place_table(db, scope).symbol_id("__init__")?;
diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index bbb07d4..c573af5 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs
@@ -205,7 +205,14 @@ Type::NominalInstance(nominal_instance) if nominal_instance.has_known_class(db, KnownClass::EllipsisType) ); - let skip_type_check = context.in_stub() && is_ellipsis; + // `auto()` values are computed at runtime by the enum metaclass, + // so we can't validate them against _value_ or __init__ at the type level. + let is_auto = matches!( + member_value_type, + Type::NominalInstance(nominal_instance) + if nominal_instance.has_known_class(db, KnownClass::Auto) + ); + let skip_type_check = (context.in_stub() && is_ellipsis) || is_auto; if !skip_type_check { if let Some(init_function) = enum_info.init_function {