[ty] support enum _value_ annotation

crates/ty_python_semantic/resources/mdtest/enums.md

@@ -100,6 +100,175 @@ class Answer(Enum):
 reveal_type(enum_members(Answer))
 ```
 
+### Declared `_value_` annotation
+
+If a `_value_` annotation is defined on an `Enum` class, all enum member values must be compatible
+with the declared type:
+
+```pyi
+from enum import Enum
+
+class Color(Enum):
+    _value_: int
+    RED = 1
+    GREEN = "green"  # error: [invalid-assignment]
+    BLUE = ...
+    YELLOW = None  # error: [invalid-assignment]
+    # In stub files, `[]` is not exempt from type checking (only `...` is).
+    PURPLE = []  # error: [invalid-assignment]
+```
+
+When `_value_` is annotated, `.value` and `._value_` are inferred as the declared type:
+
+```py
+from enum import Enum
+
+class Color2(Enum):
+    _value_: int
+    RED = 1
+    GREEN = 2
+
+reveal_type(Color2.RED.value)  # revealed: int
+reveal_type(Color2.RED._value_)  # revealed: int
+```
+
+### `_value_` annotation with `__init__`
+
+When `__init__` is defined, member values are validated by synthesizing a call to `__init__`. The
+`_value_` annotation still constrains assignments to `self._value_` inside `__init__`:
+
+```py
+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):
+        self._value_ = value  # error: [invalid-assignment]
+
+    MERCURY = (1, 3.303e23, 2.4397e6)
+    SATURN = "saturn"  # error: [invalid-assignment]
+
+reveal_type(Planet.MERCURY.value)  # revealed: str
+reveal_type(Planet.MERCURY._value_)  # revealed: str
+```
+
+### `__init__` without `_value_` annotation
+
+When `__init__` is defined but no explicit `_value_` annotation exists, member values are validated
+against the `__init__` signature. Values that are incompatible with `__init__` are flagged:
+
+```py
+from enum import Enum
+
+class Planet2(Enum):
+    def __init__(self, mass: float, radius: float):
+        self.mass = mass
+        self.radius = radius
+
+    MERCURY = (3.303e23, 2.4397e6)
+    VENUS = (4.869e24, 6.0518e6)
+    INVALID = "not a planet"  # error: [invalid-assignment]
+
+reveal_type(Planet2.MERCURY.value)  # revealed: Any
+reveal_type(Planet2.MERCURY._value_)  # revealed: Any
+```
+
+### Inherited `_value_` annotation
+
+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
+
+class Base(Enum):
+    _value_: int
+
+class Child(Base):
+    A = 1
+    B = "not an int"  # error: [invalid-assignment]
+
+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
 
 Methods, callables, descriptors (including properties), and nested classes that are defined in the
@@ -358,7 +527,8 @@ class SingleMember(StrEnum):
 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
@@ -367,8 +537,8 @@ class Answer(IntEnum):
     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:

crates/ty_python_semantic/src/types.rs

@@ -3465,7 +3465,7 @@ impl<'db> Type<'db> {
             {
                 let enum_literal = literal.as_enum().unwrap();
                 enum_metadata(db, enum_literal.enum_class(db))
-                    .and_then(|metadata| metadata.members.get(enum_literal.name(db)))
+                    .and_then(|metadata| metadata.value_type(enum_literal.name(db)))
                     .map_or_else(|| Place::Undefined, Place::bound)
                     .into()
             }
@@ -3489,10 +3489,10 @@ impl<'db> Type<'db> {
             {
                 enum_metadata(db, instance.class_literal(db))
                     .and_then(|metadata| {
-                        let (_, ty) = metadata.members.get_index(0)?;
-                        Some(Place::bound(*ty))
+                        let (name, _) = metadata.members.get_index(0)?;
+                        metadata.value_type(name)
                     })
-                    .unwrap_or_default()
+                    .map_or_else(Place::default, Place::bound)
                     .into()
             }
 

crates/ty_python_semantic/src/types/enums.rs

@@ -7,26 +7,54 @@ use crate::{
     place::{
         DefinedPlace, Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations,
     },
-    semantic_index::{place_table, use_def_map},
+    semantic_index::{place_table, scope::ScopeId, use_def_map},
     types::{
         ClassBase, ClassLiteral, DynamicType, EnumLiteralType, KnownClass, LiteralValueTypeKind,
-        MemberLookupPolicy, StaticClassLiteral, Type, TypeQualifiers,
+        MemberLookupPolicy, StaticClassLiteral, Type, TypeQualifiers, function::FunctionType,
     },
 };
 
 #[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct EnumMetadata<'db> {
     pub(crate) members: FxIndexMap<Name, Type<'db>>,
     pub(crate) aliases: FxHashMap<Name, Name>,
+
+    /// The explicit `_value_` annotation type, if declared.
+    pub(crate) value_annotation: Option<Type<'db>>,
+
+    /// The custom `__init__` function, if defined on this enum.
+    ///
+    /// When present, member values are validated by synthesizing a call to
+    /// `__init__` rather than by simple type assignability.
+    pub(crate) init_function: Option<FunctionType<'db>>,
 }
 
 impl get_size2::GetSize for EnumMetadata<'_> {}
 
-impl EnumMetadata<'_> {
+impl<'db> EnumMetadata<'db> {
     fn empty() -> Self {
         EnumMetadata {
             members: FxIndexMap::default(),
             aliases: FxHashMap::default(),
+            value_annotation: None,
+            init_function: None,
+        }
+    }
+
+    /// Returns the type of `.value`/`._value_` for a given enum member.
+    ///
+    /// Priority: explicit `_value_` annotation, then `__init__` → `Any`,
+    /// then the inferred member value type.
+    pub(crate) fn value_type(&self, member_name: &Name) -> Option<Type<'db>> {
+        if !self.members.contains_key(member_name) {
+            return None;
+        }
+        if let Some(annotation) = self.value_annotation {
+            Some(annotation)
+        } else if self.init_function.is_some() {
+            Some(Type::Dynamic(DynamicType::Any))
+        } else {
+            self.members.get(member_name).copied()
         }
     }
 
@@ -287,7 +315,93 @@ pub(crate) fn enum_metadata<'db>(
         return None;
     }
 
-    Some(EnumMetadata { members, aliases })
+    // 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| {
+            let declarations = use_def_map.end_of_scope_symbol_declarations(symbol_id);
+            place_from_declarations(db, declarations)
+                .ignore_conflicting_declarations()
+                .ignore_possibly_undefined()
+        })
+        .or_else(|| inherited_value_annotation(db, class));
+
+    // 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,
+        aliases,
+        value_annotation,
+        init_function,
+    })
+}
+
+/// 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__")?;
+    let init_type = place_from_declarations(
+        db,
+        use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id),
+    )
+    .ignore_conflicting_declarations()
+    .ignore_possibly_undefined()?;
+
+    match init_type {
+        Type::FunctionLiteral(f) => Some(f),
+        _ => None,
+    }
 }
 
 pub(crate) fn enum_member_literals<'a, 'db: 'a>(