diff --git a/crates/codegen/src/compile.rs b/crates/codegen/src/compile.rs index c531c5ff258..c58f3746ec0 100644 --- a/crates/codegen/src/compile.rs +++ b/crates/codegen/src/compile.rs @@ -29,8 +29,9 @@ use ruff_python_ast::{ InterpolatedStringElements, Keyword, MatchCase, ModExpression, ModModule, Operator, Parameters, Pattern, PatternMatchAs, PatternMatchClass, PatternMatchMapping, PatternMatchOr, PatternMatchSequence, PatternMatchSingleton, PatternMatchStar, PatternMatchValue, Singleton, - Stmt, StmtAnnAssign, StmtExpr, TString, TypeParam, TypeParamParamSpec, TypeParamTypeVar, - TypeParamTypeVarTuple, TypeParams, UnaryOp, WithItem, + Stmt, StmtAnnAssign, StmtExpr, StmtFor, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith, + TString, TypeParam, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams, + UnaryOp, WithItem, visitor::{Visitor, walk_expr}, }; use ruff_text_size::{Ranged, TextRange}; @@ -3630,23 +3631,69 @@ impl Compiler { Ok(true) } - /// Collect simple (non-conditional) annotations from module body + /// Collect simple annotations from module body in AST order (including nested blocks) /// Returns list of (name, annotation_expr) pairs + /// This must match the order that annotations are compiled to ensure + /// conditional_annotation_index stays in sync with __annotate__ enumeration. fn collect_simple_annotations(body: &[Stmt]) -> Vec<(&str, &Expr)> { - let mut annotations = Vec::new(); - for stmt in body { - if let Stmt::AnnAssign(StmtAnnAssign { - target, - annotation, - simple, - .. - }) = stmt - && *simple - && let Expr::Name(ExprName { id, .. }) = target.as_ref() - { - annotations.push((id.as_str(), annotation.as_ref())); + fn walk<'a>(stmts: &'a [Stmt], out: &mut Vec<(&'a str, &'a Expr)>) { + for stmt in stmts { + match stmt { + Stmt::AnnAssign(StmtAnnAssign { + target, + annotation, + simple, + .. + }) if *simple && matches!(target.as_ref(), Expr::Name(_)) => { + if let Expr::Name(ExprName { id, .. }) = target.as_ref() { + out.push((id.as_str(), annotation.as_ref())); + } + } + Stmt::If(StmtIf { + body, + elif_else_clauses, + .. + }) => { + walk(body, out); + for clause in elif_else_clauses { + walk(&clause.body, out); + } + } + Stmt::For(StmtFor { body, orelse, .. }) + | Stmt::While(StmtWhile { body, orelse, .. }) => { + walk(body, out); + walk(orelse, out); + } + Stmt::With(StmtWith { body, .. }) => walk(body, out), + Stmt::Try(StmtTry { + body, + handlers, + orelse, + finalbody, + .. + }) => { + walk(body, out); + for handler in handlers { + let ExceptHandler::ExceptHandler(ExceptHandlerExceptHandler { + body, + .. + }) = handler; + walk(body, out); + } + walk(orelse, out); + walk(finalbody, out); + } + Stmt::Match(StmtMatch { cases, .. }) => { + for case in cases { + walk(&case.body, out); + } + } + _ => {} + } } } + let mut annotations = Vec::new(); + walk(body, &mut annotations); annotations } diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index ef5a25f0843..49cc7168dd4 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -102,46 +102,33 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap { - Some(value.as_str().map(|s| s.to_owned()).unwrap_or_else( - |_| value.to_string_lossy().into_owned(), - )) - } - Some(other) => { - eprintln!( - "Warning: Malformed annotation for '{:?}': expected string constant at index {}, got {:?}", - param_name, val_idx, other - ); - None - } - None => { - eprintln!( - "Warning: Malformed annotation for '{:?}': constant index {} out of bounds (len={})", - param_name, - val_idx, - code.constants.len() - ); - None - } + Some(ConstantData::Str { value }) => value + .as_str() + .map(|s| s.to_owned()) + .unwrap_or_else(|_| value.to_string_lossy().into_owned()), + Some(other) => panic!( + "Unsupported annotation const for '{:?}' at idx {}: {:?}", + param_name, val_idx, other + ), + None => panic!( + "Annotation const idx out of bounds for '{:?}': {} (len={})", + param_name, + val_idx, + code.constants.len() + ), } } else { match code.names.get(val_idx) { - Some(name) => Some(name.clone()), - None => { - eprintln!( - "Warning: Malformed annotation for '{}': name index {} out of bounds (len={})", - param_name, - val_idx, - code.names.len() - ); - None - } + Some(name) => name.clone(), + None => panic!( + "Annotation name idx out of bounds for '{:?}': {} (len={})", + param_name, + val_idx, + code.names.len() + ), } }; - if let Some(type_name) = type_name { - annotations - .insert(param_name.clone(), StackValue::String(type_name)); - } + annotations.insert(param_name.clone(), StackValue::String(type_name)); } } }