From 2fd025c8667cac42f19f530f61e71497395a4c48 Mon Sep 17 00:00:00 2001 From: Simon Lamon <32477463+silamon@users.noreply.github.com> Date: Sat, 27 Dec 2025 20:12:14 +0100 Subject: [PATCH 01/12] Support enum `_value_` annotation --- .../resources/mdtest/enums.md | 32 +++++ .../ty_python_semantic/src/types/overrides.rs | 129 +++++++++++++++++- 2 files changed, 155 insertions(+), 6 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index c3b1e55c53e85..197bbb9212bca 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -100,6 +100,38 @@ 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] +``` + +Reassigning `_value_` inside `__init__` is allowed, but reassignment must still conform to the +declared `_value_` type annotation: + +```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] +``` + ### Non-member attributes with disallowed type Methods, callables, descriptors (including properties), and nested classes that are defined in the diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index eea86fe686354..ad4d0dd8ed3cd 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -12,7 +12,7 @@ use rustc_hash::FxHashSet; use crate::{ Db, lint::LintId, - place::{DefinedPlace, Place}, + place::{DefinedPlace, Place, place_from_declarations}, semantic_index::{ definition::{Definition, DefinitionKind}, place::ScopedPlaceId, @@ -27,11 +27,12 @@ use crate::{ class::{CodeGeneratorKind, FieldKind}, context::InferContext, diagnostic::{ - INVALID_DATACLASS, INVALID_EXPLICIT_OVERRIDE, INVALID_METHOD_OVERRIDE, - INVALID_NAMED_TUPLE, OVERRIDE_OF_FINAL_METHOD, OVERRIDE_OF_FINAL_VARIABLE, - report_invalid_method_override, report_overridden_final_method, - report_overridden_final_variable, + INVALID_ASSIGNMENT, INVALID_DATACLASS, INVALID_EXPLICIT_OVERRIDE, + INVALID_METHOD_OVERRIDE, INVALID_NAMED_TUPLE, OVERRIDE_OF_FINAL_METHOD, + OVERRIDE_OF_FINAL_VARIABLE, report_invalid_method_override, + report_overridden_final_method, report_overridden_final_variable, }, + enums::{EnumMetadata, enum_metadata}, function::{FunctionDecorators, FunctionType, KnownFunction}, list_members::{Member, MemberWithDefinition, all_end_of_scope_members}, }, @@ -53,6 +54,73 @@ const PROHIBITED_NAMEDTUPLE_ATTRS: &[&str] = &[ "_source", ]; +struct EnumClassInfo<'db> { + metadata: &'db EnumMetadata<'db>, + value_sunder_type: Type<'db>, +} + +impl<'db> EnumClassInfo<'db> { + fn from_class(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> Option> { + let metadata = enum_metadata(db, class.into())?; + + let scope = class.body_scope(db); + let value_sunder_symbol = place_table(db, scope).symbol_id("_value_")?; + let value_sunder_declarations = + use_def_map(db, scope).end_of_scope_symbol_declarations(value_sunder_symbol); + let value_sunder_type = place_from_declarations(db, value_sunder_declarations) + .ignore_conflicting_declarations() + .ignore_possibly_undefined()?; + + let expected_member_type = Self::extract_init_member_type(db, class, scope); + let value_sunder_type = if let Some(ty) = expected_member_type { + ty + } else { + value_sunder_type + }; + + Some(EnumClassInfo { + metadata, + value_sunder_type, + }) + } + + /// Extracts the expected enum member type from `__init__` method parameters. + fn extract_init_member_type( + db: &'db dyn Db, + _class: StaticClassLiteral<'db>, + scope: ScopeId<'db>, + ) -> Option> { + let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; + let init_declarations = + use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id); + + let place_and_qualifiers = + place_from_declarations(db, init_declarations).ignore_conflicting_declarations(); + let init_place_type = place_and_qualifiers.ignore_possibly_undefined()?; + + let Type::FunctionLiteral(func_type) = init_place_type else { + return None; + }; + + let signature = func_type.signature(db); + + let param_types: Vec> = signature + .overloads + .first()? + .parameters() + .iter() + .skip(1) // skip `self` + .map(Parameter::annotated_type) + .collect(); + + match param_types.len() { + 0 => None, + 1 => param_types.into_iter().next(), + _ => Some(Type::heterogeneous_tuple(db, param_types)), + } + } +} + // TODO: Support dynamic class literals. If we allow dynamic classes to define attributes in their // namespace dictionary, we should also check whether those attributes are valid overrides of // attributes in their superclasses. @@ -66,15 +134,24 @@ pub(super) fn check_class<'db>(context: &InferContext<'db, '_>, class: StaticCla let class_specialized = class.identity_specialization(db); let scope = class.body_scope(db); let own_class_members: FxHashSet<_> = all_end_of_scope_members(db, scope).collect(); + let enum_info = EnumClassInfo::from_class(db, class); for member in own_class_members { - check_class_declaration(context, configuration, class_specialized, scope, &member); + check_class_declaration( + context, + configuration, + enum_info.as_ref(), + class_specialized, + scope, + &member, + ); } } fn check_class_declaration<'db>( context: &InferContext<'db, '_>, configuration: OverrideRulesConfig, + enum_info: Option<&EnumClassInfo<'db>>, class: ClassType<'db>, class_scope: ScopeId<'db>, member: &MemberWithDefinition<'db>, @@ -173,6 +250,46 @@ fn check_class_declaration<'db>( Some(CodeGeneratorKind::TypedDict) | None => {} } + // Check for invalid Enum member values. + if let Some(enum_info) = enum_info { + if member.name != "_value_" + && let DefinitionKind::Assignment(_) = first_reachable_definition.kind(db) + { + let is_enum_member = enum_info.metadata.resolve_member(&member.name).is_some(); + if is_enum_member { + let member_value_type = member.ty; + + let is_ellipsis = matches!( + member_value_type, + Type::NominalInstance(nominal_instance) + if nominal_instance.has_known_class(db, KnownClass::EllipsisType) + ); + let skip_type_check = context.in_stub() && is_ellipsis; + + if !skip_type_check { + // Determine the expected type for the member + let expected_type = enum_info.value_sunder_type; + if !member_value_type.is_assignable_to(db, expected_type) { + if let Some(builder) = context.report_lint( + &INVALID_ASSIGNMENT, + first_reachable_definition.focus_range(db, context.module()), + ) { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Enum member `{}` value is not compatible with expected type", + &member.name + )); + diagnostic.info(format_args!( + "Expected type assignable to `{}`, got `{}`", + expected_type.display(db), + member_value_type.display(db) + )); + } + } + } + } + } + } + let mut subclass_overrides_superclass_declaration = false; let mut has_dynamic_superclass = false; let mut has_typeddict_in_mro = false; From 5c87d3ef7f8459d4e4b8173287fe3ddc85813417 Mon Sep 17 00:00:00 2001 From: Carl Meyer Date: Fri, 23 Jan 2026 14:17:49 -0800 Subject: [PATCH 02/12] review comments --- crates/ty_python_semantic/resources/mdtest/enums.md | 1 - crates/ty_python_semantic/src/types/overrides.rs | 9 +++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 197bbb9212bca..5ca5380c080df 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -127,7 +127,6 @@ class Planet(Enum): 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] ``` diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index ad4d0dd8ed3cd..3974236cf5e21 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -259,6 +259,9 @@ fn check_class_declaration<'db>( if is_enum_member { let member_value_type = member.ty; + // TODO ideally this would be a syntactic check that only matches on literal `...` + // in the source, rather than matching on the type. But this would require storing + // additional information in `EnumMetadata`. let is_ellipsis = matches!( member_value_type, Type::NominalInstance(nominal_instance) @@ -275,14 +278,16 @@ fn check_class_declaration<'db>( first_reachable_definition.focus_range(db, context.module()), ) { let mut diagnostic = builder.into_diagnostic(format_args!( - "Enum member `{}` value is not compatible with expected type", + "Enum member `{}` value is not assignable to expected type", &member.name )); diagnostic.info(format_args!( - "Expected type assignable to `{}`, got `{}`", + "Expected `{}`, got `{}`", expected_type.display(db), member_value_type.display(db) )); + // TODO we could also point to the source of our `_value_` type + // expectations (`_value_` annotation or `__init__` method) } } } From 581d0991443e8aedd10e558d35587b1d0b61bd2c Mon Sep 17 00:00:00 2001 From: Simon Lamon <32477463+silamon@users.noreply.github.com> Date: Sat, 24 Jan 2026 10:01:50 +0100 Subject: [PATCH 03/12] Refactor EnumClassInfo into EnumMetadata and fallback to Any if __init__ defined --- crates/ty_python_semantic/src/types/enums.rs | 62 ++++++++++++++- .../ty_python_semantic/src/types/overrides.rs | 77 ++----------------- 2 files changed, 63 insertions(+), 76 deletions(-) diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 50b80465043b4..9685ca1d8b0bc 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -7,10 +7,10 @@ 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, Parameter, StaticClassLiteral, Type, TypeQualifiers, }, }; @@ -18,15 +18,17 @@ use crate::{ pub(crate) struct EnumMetadata<'db> { pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, + pub(crate) value_sunder_type: Type<'db>, } impl get_size2::GetSize for EnumMetadata<'_> {} -impl EnumMetadata<'_> { +impl<'db> EnumMetadata<'db> { fn empty() -> Self { EnumMetadata { members: FxIndexMap::default(), aliases: FxHashMap::default(), + value_sunder_type: Type::Dynamic(DynamicType::Unknown), } } @@ -287,7 +289,59 @@ pub(crate) fn enum_metadata<'db>( return None; } - Some(EnumMetadata { members, aliases }) + // _value_ + let scope = class.body_scope(db); + let value_sunder_symbol = place_table(db, scope).symbol_id("_value_")?; + let value_sunder_declarations = + use_def_map.end_of_scope_symbol_declarations(value_sunder_symbol); + + let inferred_value_sunder_type = place_from_declarations(db, value_sunder_declarations) + .ignore_conflicting_declarations() + .ignore_possibly_undefined()?; + + let value_sunder_type = + extract_init_member_type(db, class, scope).unwrap_or(inferred_value_sunder_type); + + Some(EnumMetadata { + members, + aliases, + value_sunder_type, + }) +} + +// Extracts the expected enum member type from `__init__` method parameters. +fn extract_init_member_type<'db>( + db: &'db dyn Db, + _class: StaticClassLiteral<'db>, + scope: ScopeId<'db>, +) -> Option> { + let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; + let init_declarations = use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id); + + let place_and_qualifiers = place_from_declarations(db, init_declarations) + .ignore_conflicting_declarations() + .ignore_possibly_undefined()?; + + let Type::FunctionLiteral(func_type) = place_and_qualifiers else { + return None; + }; + + let signature = func_type.signature(db); + + let param_types: Vec> = signature + .overloads + .first()? + .parameters() + .iter() + .skip(1) // skip `self` + .map(Parameter::annotated_type) + .collect(); + + match param_types.len() { + 0 => None, + // At this moment, we just infer any as type similar to other type checkers + _ => Some(Type::Dynamic(DynamicType::Any)), + } } pub(crate) fn enum_member_literals<'a, 'db: 'a>( diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index 3974236cf5e21..27706d646d21b 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -12,7 +12,7 @@ use rustc_hash::FxHashSet; use crate::{ Db, lint::LintId, - place::{DefinedPlace, Place, place_from_declarations}, + place::{DefinedPlace, Place}, semantic_index::{ definition::{Definition, DefinitionKind}, place::ScopedPlaceId, @@ -54,73 +54,6 @@ const PROHIBITED_NAMEDTUPLE_ATTRS: &[&str] = &[ "_source", ]; -struct EnumClassInfo<'db> { - metadata: &'db EnumMetadata<'db>, - value_sunder_type: Type<'db>, -} - -impl<'db> EnumClassInfo<'db> { - fn from_class(db: &'db dyn Db, class: StaticClassLiteral<'db>) -> Option> { - let metadata = enum_metadata(db, class.into())?; - - let scope = class.body_scope(db); - let value_sunder_symbol = place_table(db, scope).symbol_id("_value_")?; - let value_sunder_declarations = - use_def_map(db, scope).end_of_scope_symbol_declarations(value_sunder_symbol); - let value_sunder_type = place_from_declarations(db, value_sunder_declarations) - .ignore_conflicting_declarations() - .ignore_possibly_undefined()?; - - let expected_member_type = Self::extract_init_member_type(db, class, scope); - let value_sunder_type = if let Some(ty) = expected_member_type { - ty - } else { - value_sunder_type - }; - - Some(EnumClassInfo { - metadata, - value_sunder_type, - }) - } - - /// Extracts the expected enum member type from `__init__` method parameters. - fn extract_init_member_type( - db: &'db dyn Db, - _class: StaticClassLiteral<'db>, - scope: ScopeId<'db>, - ) -> Option> { - let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; - let init_declarations = - use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id); - - let place_and_qualifiers = - place_from_declarations(db, init_declarations).ignore_conflicting_declarations(); - let init_place_type = place_and_qualifiers.ignore_possibly_undefined()?; - - let Type::FunctionLiteral(func_type) = init_place_type else { - return None; - }; - - let signature = func_type.signature(db); - - let param_types: Vec> = signature - .overloads - .first()? - .parameters() - .iter() - .skip(1) // skip `self` - .map(Parameter::annotated_type) - .collect(); - - match param_types.len() { - 0 => None, - 1 => param_types.into_iter().next(), - _ => Some(Type::heterogeneous_tuple(db, param_types)), - } - } -} - // TODO: Support dynamic class literals. If we allow dynamic classes to define attributes in their // namespace dictionary, we should also check whether those attributes are valid overrides of // attributes in their superclasses. @@ -134,13 +67,13 @@ pub(super) fn check_class<'db>(context: &InferContext<'db, '_>, class: StaticCla let class_specialized = class.identity_specialization(db); let scope = class.body_scope(db); let own_class_members: FxHashSet<_> = all_end_of_scope_members(db, scope).collect(); - let enum_info = EnumClassInfo::from_class(db, class); + let enum_info = enum_metadata(db, class.into()); for member in own_class_members { check_class_declaration( context, configuration, - enum_info.as_ref(), + enum_info, class_specialized, scope, &member, @@ -151,7 +84,7 @@ pub(super) fn check_class<'db>(context: &InferContext<'db, '_>, class: StaticCla fn check_class_declaration<'db>( context: &InferContext<'db, '_>, configuration: OverrideRulesConfig, - enum_info: Option<&EnumClassInfo<'db>>, + enum_info: Option<&EnumMetadata<'db>>, class: ClassType<'db>, class_scope: ScopeId<'db>, member: &MemberWithDefinition<'db>, @@ -255,7 +188,7 @@ fn check_class_declaration<'db>( if member.name != "_value_" && let DefinitionKind::Assignment(_) = first_reachable_definition.kind(db) { - let is_enum_member = enum_info.metadata.resolve_member(&member.name).is_some(); + let is_enum_member = enum_info.resolve_member(&member.name).is_some(); if is_enum_member { let member_value_type = member.ty; From da7588d074d54d6f6b698637ca12a9a922af8b82 Mon Sep 17 00:00:00 2001 From: Simon Lamon <32477463+silamon@users.noreply.github.com> Date: Sat, 7 Feb 2026 10:12:54 +0100 Subject: [PATCH 04/12] Fixes --- crates/ty_python_semantic/src/types/enums.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 9685ca1d8b0bc..b4385ff682b5b 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -23,7 +23,7 @@ pub(crate) struct EnumMetadata<'db> { impl get_size2::GetSize for EnumMetadata<'_> {} -impl<'db> EnumMetadata<'db> { +impl EnumMetadata<'_> { fn empty() -> Self { EnumMetadata { members: FxIndexMap::default(), @@ -290,8 +290,7 @@ pub(crate) fn enum_metadata<'db>( } // _value_ - let scope = class.body_scope(db); - let value_sunder_symbol = place_table(db, scope).symbol_id("_value_")?; + let value_sunder_symbol = place_table(db, scope_id).symbol_id("_value_")?; let value_sunder_declarations = use_def_map.end_of_scope_symbol_declarations(value_sunder_symbol); @@ -300,7 +299,7 @@ pub(crate) fn enum_metadata<'db>( .ignore_possibly_undefined()?; let value_sunder_type = - extract_init_member_type(db, class, scope).unwrap_or(inferred_value_sunder_type); + extract_init_member_type(db, class, scope_id).unwrap_or(inferred_value_sunder_type); Some(EnumMetadata { members, From 01d9ca8989f0f2719417b29ddcdc60fd8ed2c830 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 13 Feb 2026 20:50:45 -0500 Subject: [PATCH 05/12] Fix tests --- .../resources/mdtest/enums.md | 8 +++--- crates/ty_python_semantic/src/types/enums.rs | 25 +++++++++++-------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 5ca5380c080df..856d8cdac5a11 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -116,8 +116,8 @@ class Color(Enum): YELLOW = None # error: [invalid-assignment] ``` -Reassigning `_value_` inside `__init__` is allowed, but reassignment must still conform to the -declared `_value_` type annotation: +An explicit `_value_` annotation always takes priority, even if `__init__` is defined. If `__init__` +is defined but no `_value_` annotation exists, we fall back to `Any`: ```py from enum import Enum @@ -127,8 +127,8 @@ class Planet(Enum): 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] + MERCURY = (1, 3.303e23, 2.4397e6) # error: [invalid-assignment] + SATURN = "saturn" ``` ### Non-member attributes with disallowed type diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index b4385ff682b5b..1d305a4728245 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -289,17 +289,20 @@ pub(crate) fn enum_metadata<'db>( return None; } - // _value_ - let value_sunder_symbol = place_table(db, scope_id).symbol_id("_value_")?; - let value_sunder_declarations = - use_def_map.end_of_scope_symbol_declarations(value_sunder_symbol); - - let inferred_value_sunder_type = place_from_declarations(db, value_sunder_declarations) - .ignore_conflicting_declarations() - .ignore_possibly_undefined()?; - - let value_sunder_type = - extract_init_member_type(db, class, scope_id).unwrap_or(inferred_value_sunder_type); + // Determine the expected `_value_` type: + // (a) Always respect an explicit `_value_` annotation if present. + // (b) Otherwise, fall back to `Any` if the enum has an `__init__` method. + // (c) Otherwise, fall back to `Unknown` (no member value validation). + let value_sunder_type = 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(|| extract_init_member_type(db, class, scope_id)) + .unwrap_or(Type::Dynamic(DynamicType::Unknown)); Some(EnumMetadata { members, From c265c6f561c6ae5a880be13f44e2a88269722080 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 13 Feb 2026 20:59:43 -0500 Subject: [PATCH 06/12] Nits --- crates/ty_python_semantic/resources/mdtest/enums.md | 1 + crates/ty_python_semantic/src/types/enums.rs | 8 ++------ crates/ty_python_semantic/src/types/overrides.rs | 5 ++++- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 856d8cdac5a11..16ca405309938 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -127,6 +127,7 @@ class Planet(Enum): def __init__(self, value: int, mass: float, radius: float): self._value_ = value # error: [invalid-assignment] + MERCURY = (1, 3.303e23, 2.4397e6) # error: [invalid-assignment] SATURN = "saturn" ``` diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 1d305a4728245..1a3069e09d0fb 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -301,7 +301,7 @@ pub(crate) fn enum_metadata<'db>( .ignore_conflicting_declarations() .ignore_possibly_undefined() }) - .or_else(|| extract_init_member_type(db, class, scope_id)) + .or_else(|| extract_init_member_type(db, scope_id)) .unwrap_or(Type::Dynamic(DynamicType::Unknown)); Some(EnumMetadata { @@ -312,11 +312,7 @@ pub(crate) fn enum_metadata<'db>( } // Extracts the expected enum member type from `__init__` method parameters. -fn extract_init_member_type<'db>( - db: &'db dyn Db, - _class: StaticClassLiteral<'db>, - scope: ScopeId<'db>, -) -> Option> { +fn extract_init_member_type<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option> { let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; let init_declarations = use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id); diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index 27706d646d21b..b259d3754f2e3 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -186,7 +186,10 @@ fn check_class_declaration<'db>( // Check for invalid Enum member values. if let Some(enum_info) = enum_info { if member.name != "_value_" - && let DefinitionKind::Assignment(_) = first_reachable_definition.kind(db) + && matches!( + first_reachable_definition.kind(db), + DefinitionKind::Assignment(_) + ) { let is_enum_member = enum_info.resolve_member(&member.name).is_some(); if is_enum_member { From ba7371d32d494151a17aeedd0af88e7afc7d7378 Mon Sep 17 00:00:00 2001 From: Simon Lamon <32477463+silamon@users.noreply.github.com> Date: Sat, 14 Feb 2026 10:04:50 +0100 Subject: [PATCH 07/12] Fix `mercury` example --- .../resources/mdtest/enums.md | 8 ++--- crates/ty_python_semantic/src/types/enums.rs | 34 +++++++++++-------- 2 files changed, 24 insertions(+), 18 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 16ca405309938..6f4767ee54456 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -116,8 +116,8 @@ class Color(Enum): YELLOW = None # error: [invalid-assignment] ``` -An explicit `_value_` annotation always takes priority, even if `__init__` is defined. If `__init__` -is defined but no `_value_` annotation exists, we fall back to `Any`: +If `__init__` is defined, any explicit `_value_` annotation is ignored and the deferred `_value_` +type comes from the initializer signature. ```py from enum import Enum @@ -128,8 +128,8 @@ class Planet(Enum): def __init__(self, value: int, mass: float, radius: float): self._value_ = value # error: [invalid-assignment] - MERCURY = (1, 3.303e23, 2.4397e6) # error: [invalid-assignment] - SATURN = "saturn" + MERCURY = (1, 3.303e23, 2.4397e6) + SATURN = "saturn" # error: [invalid-assignment] ``` ### Non-member attributes with disallowed type diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 1a3069e09d0fb..1d5dad70dee9d 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -10,7 +10,7 @@ use crate::{ semantic_index::{place_table, scope::ScopeId, use_def_map}, types::{ ClassBase, ClassLiteral, DynamicType, EnumLiteralType, KnownClass, LiteralValueTypeKind, - MemberLookupPolicy, Parameter, StaticClassLiteral, Type, TypeQualifiers, + MemberLookupPolicy, Parameter, StaticClassLiteral, Type, TypeQualifiers, tuple::TupleType, }, }; @@ -290,18 +290,20 @@ pub(crate) fn enum_metadata<'db>( } // Determine the expected `_value_` type: - // (a) Always respect an explicit `_value_` annotation if present. - // (b) Otherwise, fall back to `Any` if the enum has an `__init__` method. - // (c) Otherwise, fall back to `Unknown` (no member value validation). - let value_sunder_type = 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() + // (a) If `__init__` is defined, the deferred `_value_` type comes from the initializer signature. + // (b) Otherwise, respect an explicit `_value_` annotation if present. + // (c) If neither exists, fall back to `Unknown` (no member value validation). + let value_sunder_type = extract_init_member_type(db, scope_id) + .or_else(|| { + 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(|| extract_init_member_type(db, scope_id)) .unwrap_or(Type::Dynamic(DynamicType::Unknown)); Some(EnumMetadata { @@ -337,8 +339,12 @@ fn extract_init_member_type<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option match param_types.len() { 0 => None, - // At this moment, we just infer any as type similar to other type checkers - _ => Some(Type::Dynamic(DynamicType::Any)), + // single-argument `__init__` – the value type is just that parameter. + 1 => Some(param_types.into_iter().next().unwrap()), + // multiple parameters – the member value is constructed from a tuple of + // all arguments, so that `M = (a, b, c)` is valid when the signature is + // `def __init__(self, a: A, b: B, c: C)`. + _ => TupleType::heterogeneous(db, param_types).map(|tuple| Type::tuple(Some(tuple))), } } From 45de0bf2a207e80577cc24fd6d12efd27ec20327 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 14 Feb 2026 18:39:40 -0500 Subject: [PATCH 08/12] Fall back to Any --- .../resources/mdtest/enums.md | 46 +++++++- crates/ty_python_semantic/src/types.rs | 8 +- crates/ty_python_semantic/src/types/enums.rs | 106 ++++++++++-------- 3 files changed, 104 insertions(+), 56 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 6f4767ee54456..697fafabaac33 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -116,8 +116,25 @@ class Color(Enum): YELLOW = None # error: [invalid-assignment] ``` -If `__init__` is defined, any explicit `_value_` annotation is ignored and the deferred `_value_` -type comes from the initializer signature. +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 passed through `__init__` rather than directly +assigned to `_value_`, so we fall back to `Any` for member value validation. The `_value_` +annotation still constrains assignments to `self._value_` inside `__init__`: ```py from enum import Enum @@ -129,7 +146,30 @@ class Planet(Enum): self._value_ = value # error: [invalid-assignment] MERCURY = (1, 3.303e23, 2.4397e6) - SATURN = "saturn" # error: [invalid-assignment] + SATURN = "saturn" + +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, we also fall back to `Any` +for member value validation: + +```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) + +reveal_type(Planet2.MERCURY.value) # revealed: Any +reveal_type(Planet2.MERCURY._value_) # revealed: Any ``` ### Non-member attributes with disallowed type diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index a4d17e190a163..d621755c3e82f 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/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() } diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 1d5dad70dee9d..490c2a795ee83 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -10,7 +10,7 @@ use crate::{ semantic_index::{place_table, scope::ScopeId, use_def_map}, types::{ ClassBase, ClassLiteral, DynamicType, EnumLiteralType, KnownClass, LiteralValueTypeKind, - MemberLookupPolicy, Parameter, StaticClassLiteral, Type, TypeQualifiers, tuple::TupleType, + MemberLookupPolicy, StaticClassLiteral, Type, TypeQualifiers, }, }; @@ -18,20 +18,46 @@ use crate::{ pub(crate) struct EnumMetadata<'db> { pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, + + /// The type used for *validating* member value assignments. + /// + /// Priority: `__init__` → `Any`, else `_value_` annotation, else `Unknown`. pub(crate) value_sunder_type: Type<'db>, + + /// The explicit `_value_` annotation type, if declared. + /// + /// This is kept separate from `value_sunder_type` because `.value` access + /// always prefers the `_value_` annotation, even when `__init__` exists. + value_annotation: Option>, } impl get_size2::GetSize for EnumMetadata<'_> {} -impl EnumMetadata<'_> { +impl<'db> EnumMetadata<'db> { fn empty() -> Self { EnumMetadata { members: FxIndexMap::default(), aliases: FxHashMap::default(), value_sunder_type: Type::Dynamic(DynamicType::Unknown), + value_annotation: 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> { + let inferred = self.members.get(member_name).copied()?; + if let Some(annotation) = self.value_annotation { + return Some(annotation); + } + Some(match self.value_sunder_type { + Type::Dynamic(DynamicType::Unknown) => inferred, + declared => declared, + }) + } + pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> { if self.members.contains_key(name) { Some(name) @@ -289,63 +315,45 @@ pub(crate) fn enum_metadata<'db>( return None; } - // Determine the expected `_value_` type: - // (a) If `__init__` is defined, the deferred `_value_` type comes from the initializer signature. - // (b) Otherwise, respect an explicit `_value_` annotation if present. - // (c) If neither exists, fall back to `Unknown` (no member value validation). - let value_sunder_type = extract_init_member_type(db, scope_id) - .or_else(|| { - 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() - }) - }) + // Look up an explicit `_value_` annotation, if present. + 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() + }); + + // Determine the expected type for member value validation: + // (a) If `__init__` is defined, fall back to `Any` (member values are passed + // through `__init__`, not directly assigned to `_value_`). + // (b) Otherwise, use an explicit `_value_` annotation if present. + // (c) Otherwise, fall back to `Unknown` (no member value validation). + let value_sunder_type = has_custom_init(db, scope_id) + .or(value_annotation) .unwrap_or(Type::Dynamic(DynamicType::Unknown)); Some(EnumMetadata { members, aliases, value_sunder_type, + value_annotation, }) } -// Extracts the expected enum member type from `__init__` method parameters. -fn extract_init_member_type<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option> { +/// If the enum defines a custom `__init__`, member values are passed through it +/// rather than being assigned directly to `_value_`, so we fall back to `Any`. +fn has_custom_init<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option> { let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; - let init_declarations = use_def_map(db, scope).end_of_scope_symbol_declarations(init_symbol_id); - - let place_and_qualifiers = place_from_declarations(db, init_declarations) - .ignore_conflicting_declarations() - .ignore_possibly_undefined()?; - - let Type::FunctionLiteral(func_type) = place_and_qualifiers else { - return None; - }; - - let signature = func_type.signature(db); - - let param_types: Vec> = signature - .overloads - .first()? - .parameters() - .iter() - .skip(1) // skip `self` - .map(Parameter::annotated_type) - .collect(); - - match param_types.len() { - 0 => None, - // single-argument `__init__` – the value type is just that parameter. - 1 => Some(param_types.into_iter().next().unwrap()), - // multiple parameters – the member value is constructed from a tuple of - // all arguments, so that `M = (a, b, c)` is valid when the signature is - // `def __init__(self, a: A, b: B, c: C)`. - _ => TupleType::heterogeneous(db, param_types).map(|tuple| Type::tuple(Some(tuple))), - } + 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()?; + + matches!(init_type, Type::FunctionLiteral(_)).then_some(Type::Dynamic(DynamicType::Any)) } pub(crate) fn enum_member_literals<'a, 'db: 'a>( From de50a9f64fd4f1e20f3ef92deccfe7f7368539a9 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 14 Feb 2026 18:46:07 -0500 Subject: [PATCH 09/12] More tests --- .../resources/mdtest/enums.md | 20 +++++++++++++++++++ crates/ty_python_semantic/src/types/enums.rs | 13 ++++++------ 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 697fafabaac33..67fbdb67c0ff0 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -114,6 +114,8 @@ class Color(Enum): 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: @@ -172,6 +174,24 @@ 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 not inherited by subclasses for the purpose of member +value validation: + +```py +from enum import Enum + +class Base(Enum): + _value_: int + +class Child(Base): + A = 1 + B = "not checked against int" + +reveal_type(Child.A.value) # revealed: Literal[1] +``` + ### Non-member attributes with disallowed type Methods, callables, descriptors (including properties), and nested classes that are defined in the diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 490c2a795ee83..9786e9e9f0c0f 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -48,14 +48,15 @@ impl<'db> EnumMetadata<'db> { /// Priority: explicit `_value_` annotation, then `__init__` → `Any`, /// then the inferred member value type. pub(crate) fn value_type(&self, member_name: &Name) -> Option> { - let inferred = self.members.get(member_name).copied()?; if let Some(annotation) = self.value_annotation { - return Some(annotation); + // Check the member exists, but use the declared annotation type. + self.members.contains_key(member_name).then_some(annotation) + } else { + match self.value_sunder_type { + Type::Dynamic(DynamicType::Unknown) => self.members.get(member_name).copied(), + declared => self.members.contains_key(member_name).then_some(declared), + } } - Some(match self.value_sunder_type { - Type::Dynamic(DynamicType::Unknown) => inferred, - declared => declared, - }) } pub(crate) fn resolve_member<'a>(&'a self, name: &'a Name) -> Option<&'a Name> { From 6ddf331706344fe4a2b08b4fa96680922ad32a6d Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 16 Feb 2026 12:56:39 -0500 Subject: [PATCH 10/12] Retain validation --- .../resources/mdtest/enums.md | 12 +- crates/ty_python_semantic/src/types/enums.rs | 38 +++++-- .../ty_python_semantic/src/types/overrides.rs | 106 ++++++++++++++---- 3 files changed, 119 insertions(+), 37 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 67fbdb67c0ff0..3ca761fae8021 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -134,9 +134,8 @@ reveal_type(Color2.RED._value_) # revealed: int ### `_value_` annotation with `__init__` -When `__init__` is defined, member values are passed through `__init__` rather than directly -assigned to `_value_`, so we fall back to `Any` for member value validation. The `_value_` -annotation still constrains assignments to `self._value_` inside `__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 @@ -148,7 +147,7 @@ class Planet(Enum): self._value_ = value # error: [invalid-assignment] MERCURY = (1, 3.303e23, 2.4397e6) - SATURN = "saturn" + SATURN = "saturn" # error: [invalid-assignment] reveal_type(Planet.MERCURY.value) # revealed: str reveal_type(Planet.MERCURY._value_) # revealed: str @@ -156,8 +155,8 @@ reveal_type(Planet.MERCURY._value_) # revealed: str ### `__init__` without `_value_` annotation -When `__init__` is defined but no explicit `_value_` annotation exists, we also fall back to `Any` -for member value validation: +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 @@ -169,6 +168,7 @@ class Planet2(Enum): 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 diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 9786e9e9f0c0f..43aaf89adfdea 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -10,7 +10,7 @@ use crate::{ 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, }, }; @@ -19,9 +19,10 @@ pub(crate) struct EnumMetadata<'db> { pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, - /// The type used for *validating* member value assignments. + /// The type used for `.value` access when no `_value_` annotation exists. /// - /// Priority: `__init__` → `Any`, else `_value_` annotation, else `Unknown`. + /// When `__init__` is defined, this is `Any` (since values are processed + /// through `__init__`). Otherwise, falls back to `Unknown` (use inferred type). pub(crate) value_sunder_type: Type<'db>, /// The explicit `_value_` annotation type, if declared. @@ -29,6 +30,12 @@ pub(crate) struct EnumMetadata<'db> { /// This is kept separate from `value_sunder_type` because `.value` access /// always prefers the `_value_` annotation, even when `__init__` exists. value_annotation: Option>, + + /// 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>, } impl get_size2::GetSize for EnumMetadata<'_> {} @@ -40,6 +47,7 @@ impl<'db> EnumMetadata<'db> { aliases: FxHashMap::default(), value_sunder_type: Type::Dynamic(DynamicType::Unknown), value_annotation: None, + init_function: None, } } @@ -326,12 +334,15 @@ pub(crate) fn enum_metadata<'db>( .ignore_possibly_undefined() }); - // Determine the expected type for member value validation: - // (a) If `__init__` is defined, fall back to `Any` (member values are passed - // through `__init__`, not directly assigned to `_value_`). + let init_function = custom_init(db, scope_id); + + // Determine the type for `.value` access and simple member value validation: + // (a) If `__init__` is defined, fall back to `Any` for `.value` access + // (member value validation is handled separately via call synthesis). // (b) Otherwise, use an explicit `_value_` annotation if present. - // (c) Otherwise, fall back to `Unknown` (no member value validation). - let value_sunder_type = has_custom_init(db, scope_id) + // (c) Otherwise, fall back to `Unknown` (use inferred member value type). + let value_sunder_type = init_function + .map(|_| Type::Dynamic(DynamicType::Any)) .or(value_annotation) .unwrap_or(Type::Dynamic(DynamicType::Unknown)); @@ -340,12 +351,12 @@ pub(crate) fn enum_metadata<'db>( aliases, value_sunder_type, value_annotation, + init_function, }) } -/// If the enum defines a custom `__init__`, member values are passed through it -/// rather than being assigned directly to `_value_`, so we fall back to `Any`. -fn has_custom_init<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option> { +/// 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> { let init_symbol_id = place_table(db, scope).symbol_id("__init__")?; let init_type = place_from_declarations( db, @@ -354,7 +365,10 @@ fn has_custom_init<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Option Some(f), + _ => None, + } } pub(crate) fn enum_member_literals<'a, 'db: 'a>( diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index b259d3754f2e3..2cf51586a0ed3 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -23,7 +23,8 @@ use crate::{ }, types::{ CallableType, ClassBase, ClassType, KnownClass, Parameter, Parameters, Signature, - StaticClassLiteral, Type, TypeQualifiers, + StaticClassLiteral, Type, TypeContext, TypeQualifiers, + call::CallArguments, class::{CodeGeneratorKind, FieldKind}, context::InferContext, diagnostic::{ @@ -35,6 +36,7 @@ use crate::{ enums::{EnumMetadata, enum_metadata}, function::{FunctionDecorators, FunctionType, KnownFunction}, list_members::{Member, MemberWithDefinition, all_end_of_scope_members}, + tuple::Tuple, }, }; @@ -206,24 +208,32 @@ fn check_class_declaration<'db>( let skip_type_check = context.in_stub() && is_ellipsis; if !skip_type_check { - // Determine the expected type for the member - let expected_type = enum_info.value_sunder_type; - if !member_value_type.is_assignable_to(db, expected_type) { - if let Some(builder) = context.report_lint( - &INVALID_ASSIGNMENT, - first_reachable_definition.focus_range(db, context.module()), - ) { - let mut diagnostic = builder.into_diagnostic(format_args!( - "Enum member `{}` value is not assignable to expected type", - &member.name - )); - diagnostic.info(format_args!( - "Expected `{}`, got `{}`", - expected_type.display(db), - member_value_type.display(db) - )); - // TODO we could also point to the source of our `_value_` type - // expectations (`_value_` annotation or `__init__` method) + if let Some(init_function) = enum_info.init_function { + check_enum_member_against_init( + context, + init_function, + instance_of_class, + member_value_type, + &member.name, + *first_reachable_definition, + ); + } else { + let expected_type = enum_info.value_sunder_type; + if !member_value_type.is_assignable_to(db, expected_type) { + if let Some(builder) = context.report_lint( + &INVALID_ASSIGNMENT, + first_reachable_definition.focus_range(db, context.module()), + ) { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Enum member `{}` value is not assignable to expected type", + &member.name + )); + diagnostic.info(format_args!( + "Expected `{}`, got `{}`", + expected_type.display(db), + member_value_type.display(db) + )); + } } } } @@ -711,3 +721,61 @@ fn check_post_init_signature<'db>( as positional-only parameters", ); } + +/// Validates an enum member value against the enum's `__init__` signature. +/// +/// The enum metaclass unpacks tuple values as positional arguments to `__init__`, +/// and passes non-tuple values as a single argument. This function synthesizes +/// a call to `__init__` with the appropriate arguments and reports a diagnostic +/// if the call would fail. +fn check_enum_member_against_init<'db>( + context: &InferContext<'db, '_>, + init_function: FunctionType<'db>, + self_type: Type<'db>, + member_value_type: Type<'db>, + member_name: &Name, + definition: Definition<'db>, +) { + let db = context.db(); + + // The enum metaclass unpacks tuple values as positional args: + // MEMBER = (a, b, c) → __init__(self, a, b, c) + // MEMBER = x → __init__(self, x) + let args: Vec> = if let Type::NominalInstance(instance) = member_value_type { + if let Some(spec) = instance.tuple_spec(db) { + if let Tuple::Fixed(fixed) = &*spec { + fixed.all_elements().to_vec() + } else { + // Variable-length tuples: can't determine exact args, skip validation. + return; + } + } else { + vec![member_value_type] + } + } else { + vec![member_value_type] + }; + + let call_args = CallArguments::positional(args); + let call_args = call_args.with_self(Some(self_type)); + + let result = Type::FunctionLiteral(init_function) + .bindings(db) + .match_parameters(db, &call_args) + .check_types(db, &call_args, TypeContext::default(), &[]); + + if result.is_err() { + if let Some(builder) = context.report_lint( + &INVALID_ASSIGNMENT, + definition.focus_range(db, context.module()), + ) { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Enum member `{member_name}` is incompatible with `__init__`", + )); + diagnostic.info(format_args!( + "Expected compatible arguments for `{}`", + Type::FunctionLiteral(init_function).display(db), + )); + } + } +} From ad839d1640951c49d4ec7c62e32c6eec9e3f5cd0 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 20 Feb 2026 21:13:20 -0500 Subject: [PATCH 11/12] Remove value_sunder_type --- crates/ty_python_semantic/src/types/enums.rs | 36 +++++-------------- .../ty_python_semantic/src/types/overrides.rs | 3 +- 2 files changed, 9 insertions(+), 30 deletions(-) diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 43aaf89adfdea..2ecf7e1097596 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -19,17 +19,8 @@ pub(crate) struct EnumMetadata<'db> { pub(crate) members: FxIndexMap>, pub(crate) aliases: FxHashMap, - /// The type used for `.value` access when no `_value_` annotation exists. - /// - /// When `__init__` is defined, this is `Any` (since values are processed - /// through `__init__`). Otherwise, falls back to `Unknown` (use inferred type). - pub(crate) value_sunder_type: Type<'db>, - /// The explicit `_value_` annotation type, if declared. - /// - /// This is kept separate from `value_sunder_type` because `.value` access - /// always prefers the `_value_` annotation, even when `__init__` exists. - value_annotation: Option>, + pub(crate) value_annotation: Option>, /// The custom `__init__` function, if defined on this enum. /// @@ -45,7 +36,6 @@ impl<'db> EnumMetadata<'db> { EnumMetadata { members: FxIndexMap::default(), aliases: FxHashMap::default(), - value_sunder_type: Type::Dynamic(DynamicType::Unknown), value_annotation: None, init_function: None, } @@ -56,14 +46,15 @@ impl<'db> EnumMetadata<'db> { /// Priority: explicit `_value_` annotation, then `__init__` → `Any`, /// then the inferred member value type. pub(crate) fn value_type(&self, member_name: &Name) -> Option> { + if !self.members.contains_key(member_name) { + return None; + } if let Some(annotation) = self.value_annotation { - // Check the member exists, but use the declared annotation type. - self.members.contains_key(member_name).then_some(annotation) + Some(annotation) + } else if self.init_function.is_some() { + Some(Type::Dynamic(DynamicType::Any)) } else { - match self.value_sunder_type { - Type::Dynamic(DynamicType::Unknown) => self.members.get(member_name).copied(), - declared => self.members.contains_key(member_name).then_some(declared), - } + self.members.get(member_name).copied() } } @@ -336,20 +327,9 @@ pub(crate) fn enum_metadata<'db>( let init_function = custom_init(db, scope_id); - // Determine the type for `.value` access and simple member value validation: - // (a) If `__init__` is defined, fall back to `Any` for `.value` access - // (member value validation is handled separately via call synthesis). - // (b) Otherwise, use an explicit `_value_` annotation if present. - // (c) Otherwise, fall back to `Unknown` (use inferred member value type). - let value_sunder_type = init_function - .map(|_| Type::Dynamic(DynamicType::Any)) - .or(value_annotation) - .unwrap_or(Type::Dynamic(DynamicType::Unknown)); - Some(EnumMetadata { members, aliases, - value_sunder_type, value_annotation, init_function, }) diff --git a/crates/ty_python_semantic/src/types/overrides.rs b/crates/ty_python_semantic/src/types/overrides.rs index 2cf51586a0ed3..bbb07d40ac5c7 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -217,8 +217,7 @@ fn check_class_declaration<'db>( &member.name, *first_reachable_definition, ); - } else { - let expected_type = enum_info.value_sunder_type; + } else if let Some(expected_type) = enum_info.value_annotation { if !member_value_type.is_assignable_to(db, expected_type) { if let Some(builder) = context.report_lint( &INVALID_ASSIGNMENT, From 7590221c99ef01312afcd9c2531a5fe31db1f1fa Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 23 Feb 2026 11:55:59 -0500 Subject: [PATCH 12/12] Validate against MRO --- .../resources/mdtest/enums.md | 92 +++++++++++++++++-- crates/ty_python_semantic/src/types/enums.rs | 59 +++++++++++- .../ty_python_semantic/src/types/overrides.rs | 9 +- 3 files changed, 149 insertions(+), 11 deletions(-) diff --git a/crates/ty_python_semantic/resources/mdtest/enums.md b/crates/ty_python_semantic/resources/mdtest/enums.md index 3ca761fae8021..8f7088c8eaf4a 100644 --- a/crates/ty_python_semantic/resources/mdtest/enums.md +++ b/crates/ty_python_semantic/resources/mdtest/enums.md @@ -140,6 +140,26 @@ When `__init__` is defined, member values are validated by synthesizing a call t ```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 @@ -176,8 +196,8 @@ reveal_type(Planet2.MERCURY._value_) # revealed: Any ### 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 Base(Enum): class Child(Base): A = 1 - B = "not checked against int" + 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 -reveal_type(Child.A.value) # revealed: Literal[1] +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 @@ 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 @@ -459,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: diff --git a/crates/ty_python_semantic/src/types/enums.rs b/crates/ty_python_semantic/src/types/enums.rs index 2ecf7e1097596..0cb4eb7632091 100644 --- a/crates/ty_python_semantic/src/types/enums.rs +++ b/crates/ty_python_semantic/src/types/enums.rs @@ -315,7 +315,8 @@ pub(crate) fn enum_metadata<'db>( 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 @@ pub(crate) fn enum_metadata<'db>( 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 @@ pub(crate) fn enum_metadata<'db>( }) } +/// 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> + '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> { + 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> { + 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> { 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 bbb07d40ac5c7..c573af5c3dd2d 100644 --- a/crates/ty_python_semantic/src/types/overrides.rs +++ b/crates/ty_python_semantic/src/types/overrides.rs @@ -205,7 +205,14 @@ fn check_class_declaration<'db>( 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 {