8000 support annotations in nested block by youknowone · Pull Request #6767 · RustPython/RustPython · GitHub
[go: up one dir, main page]
Skip to content

support annotations in nested block #6767

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Jan 18, 2026
Merged

support annotations in nested block #6767

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 62 additions & 15 deletions crates/codegen/src/compile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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
}

Expand Down
57 changes: 22 additions & 35 deletions crates/jit/tests/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,46 +102,33 @@ fn extract_annotations_from_annotate_code(code: &CodeObject) -> HashMap<Wtf8Buf,
// Value can be a name (type ref) or a const string (forward ref)
let type_name = if val_is_const {
match code.constants.get(val_idx) {
Some(ConstantData::Str { value }) => {
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));
}
}
}
Expand Down
Loading
0