Skip to content

Commit eb50246

Browse files
youknowoneyouknowone
authored
fix annotation (#6767)
1 parent 7c7b824 commit eb50246

File tree

2 files changed

+84
-50
lines changed

2 files changed

+84
-50
lines changed

‎crates/codegen/src/compile.rs‎

Lines changed: 62 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,9 @@ use ruff_python_ast::{
2929
InterpolatedStringElements, Keyword, MatchCase, ModExpression, ModModule, Operator, Parameters,
3030
Pattern, PatternMatchAs, PatternMatchClass, PatternMatchMapping, PatternMatchOr,
3131
PatternMatchSequence, PatternMatchSingleton, PatternMatchStar, PatternMatchValue, Singleton,
32-
Stmt, StmtAnnAssign, StmtExpr, TString, TypeParam, TypeParamParamSpec, TypeParamTypeVar,
33-
TypeParamTypeVarTuple, TypeParams, UnaryOp, WithItem,
32+
Stmt, StmtAnnAssign, StmtExpr, StmtFor, StmtIf, StmtMatch, StmtTry, StmtWhile, StmtWith,
33+
TString, TypeParam, TypeParamParamSpec, TypeParamTypeVar, TypeParamTypeVarTuple, TypeParams,
34+
UnaryOp, WithItem,
3435
visitor::{Visitor, walk_expr},
3536
};
3637
use ruff_text_size::{Ranged, TextRange};
@@ -3630,23 +3631,69 @@ impl Compiler {
36303631
Ok(true)
36313632
}
36323633

3633-
/// Collect simple (non-conditional) annotations from module body
3634+
/// Collect simple annotations from module body in AST order (including nested blocks)
36343635
/// Returns list of (name, annotation_expr) pairs
3636+
/// This must match the order that annotations are compiled to ensure
3637+
/// conditional_annotation_index stays in sync with __annotate__ enumeration.
36353638
fn collect_simple_annotations(body: &[Stmt]) -> Vec<(&str, &Expr)> {
3636-
let mut annotations = Vec::new();
3637-
for stmt in body {
3638-
if let Stmt::AnnAssign(StmtAnnAssign {
3639-
target,
3640-
annotation,
3641-
simple,
3642-
..
3643-
}) = stmt
3644-
&& *simple
3645-
&& let Expr::Name(ExprName { id, .. }) = target.as_ref()
3646-
{
3647-
annotations.push((id.as_str(), annotation.as_ref()));
3639+
fn walk<'a>(stmts: &'a [Stmt], out: &mut Vec<(&'a str, &'a Expr)>) {
3640+
for stmt in stmts {
3641+
match stmt {
3642+
Stmt::AnnAssign(StmtAnnAssign {
3643+
target,
3644+
annotation,
3645+
simple,
3646+
..
3647+
}) if *simple && matches!(target.as_ref(), Expr::Name(_)) => {
3648+
if let Expr::Name(ExprName { id, .. }) = target.as_ref() {
3649+
out.push((id.as_str(), annotation.as_ref()));
3650+
}
3651+
}
3652+
Stmt::If(StmtIf {
3653+
body,
3654+
elif_else_clauses,
3655+
..
3656+
}) => {
3657+
walk(body, out);
3658+
for clause in elif_else_clauses {
3659+
walk(&clause.body, out);
3660+
}
3661+
}
3662+
Stmt::For(StmtFor { body, orelse, .. })
3663+
| Stmt::While(StmtWhile { body, orelse, .. }) => {
3664+
walk(body, out);
3665+
walk(orelse, out);
3666+
}
3667+
Stmt::With(StmtWith { body, .. }) => walk(body, out),
3668+
Stmt::Try(StmtTry {
3669+
body,
3670+
handlers,
3671+
orelse,
3672+
finalbody,
3673+
..
3674+
}) => {
3675+
walk(body, out);
3676+
for handler in handlers {
3677+
let ExceptHandler::ExceptHandler(ExceptHandlerExceptHandler {
3678+
body,
3679+
..
3680+
}) = handler;
3681+
walk(body, out);
3682+
}
3683+
walk(orelse, out);
3684+
walk(finalbody, out);
3685+
}
3686+
Stmt::Match(StmtMatch { cases, .. }) => {
3687+
for case in cases {
3688+
walk(&case.body, out);
3689+
}
3690+
}
3691+
_ => {}
3692+
}
36483693
}
36493694
}
3695+
let mut annotations = Vec::new();
3696+
walk(body, &mut annotations);
36503697
annotations
36513698
}
36523699

‎crates/jit/tests/common.rs‎

Lines changed: 22 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -102,46 +102,33 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap<Wtf8Buf,
102102
// Value can be a name (type ref) or a const string (forward ref)
103103
let type_name = if val_is_const {
104104
match code.constants.get(val_idx) {
105-
Some(ConstantData::Str { value }) => {
106-
Some(value.as_str().map(|s| s.to_owned()).unwrap_or_else(
107-
|_| value.to_string_lossy().into_owned(),
108-
))
109-
}
110-
Some(other) => {
111-
eprintln!(
112-
"Warning: Malformed annotation for '{:?}': expected string constant at index {}, got {:?}",
113-
param_name, val_idx, other
114-
);
115-
None
116-
}
117-
None => {
118-
eprintln!(
119-
"Warning: Malformed annotation for '{:?}': constant index {} out of bounds (len={})",
120-
param_name,
121-
val_idx,
122-
code.constants.len()
123-
);
124-
None
125-
}
105+
Some(ConstantData::Str { value }) => value
106+
.as_str()
107+
.map(|s| s.to_owned())
108+
.unwrap_or_else(|_| value.to_string_lossy().into_owned()),
109+
Some(other) => panic!(
110+
"Unsupported annotation const for '{:?}' at idx {}: {:?}",
111+
param_name, val_idx, other
112+
),
113+
None => panic!(
114+
"Annotation const idx out of bounds for '{:?}': {} (len={})",
115+
param_name,
116+
val_idx,
117+
code.constants.len()
118+
),
126119
}
127120
} else {
128121
match code.names.get(val_idx) {
129-
Some(name) => Some(name.clone()),
130-
None => {
131-
eprintln!(
132-
"Warning: Malformed annotation for '{}': name index {} out of bounds (len={})",
133-
param_name,
134-
val_idx,
135-
code.names.len()
136-
);
137-
None
138-
}
122+
Some(name) => name.clone(),
123+
None => panic!(
124+
"Annotation name idx out of bounds for '{:?}': {} (len={})",
125+
param_name,
126+
val_idx,
127+
code.names.len()
128+
),
139129
}
140130
};
141-
if let Some(type_name) = type_name {
142-
annotations
143-
.insert(param_name.clone(), StackValue::String(type_name));
144-
}
131+
annotations.insert(param_name.clone(), StackValue::String(type_name));
145132
}
146133
}
147134
}

0 commit comments

Comments
 (0)