Skip to content

Commit 43eb720

Browse files
youknowoneyouknowone
committed
Fix AST field defaults and compile() type check
- Extract empty_arguments_object helper from expression.rs - Fix LIST_FIELDS ambiguity: "args" and "body" have different ASDL types per node (e.g. Lambda.args is `arguments`, not `expr*`; Lambda.body is `expr`, not `stmt*`) - Replace class name string comparison in compile() with fast_isinstance to accept AST subclasses
1 parent aff89c2 commit 43eb720

File tree

5 files changed

+88
-49
lines changed

5 files changed

+88
-49
lines changed

‎crates/vm/src/protocol/object.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -705,7 +705,9 @@ impl PyObject {
705705
if let Some(class_getitem) =
706706
vm.get_attribute_opt(self.to_owned(), identifier!(vm, __class_getitem__))?
707707
{
708-
return class_getitem.call((needle,), vm);
708+
if !vm.is_none(&class_getitem) {
709+
return class_getitem.call((needle,), vm);
710+
}
709711
}
710712
return Err(vm.new_type_error(format!(
711713
"type '{}' is not subscriptable",

‎crates/vm/src/stdlib/ast.rs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,42 @@ fn node_add_location(
265265
.unwrap();
266266
}
267267

268+
/// Return the expected AST mod type class for a compile() mode string.
269+
pub(crate) fn mode_type_and_name(
270+
ctx: &Context,
271+
mode: &str,
272+
) -> Option<(PyRef<PyType>, &'static str)> {
273+
match mode {
274+
"exec" => Some((pyast::NodeModModule::make_class(ctx), "Module")),
275+
"eval" => Some((pyast::NodeModExpression::make_class(ctx), "Expression")),
276+
"single" => Some((pyast::NodeModInteractive::make_class(ctx), "Interactive")),
277+
"func_type" => Some((pyast::NodeModFunctionType::make_class(ctx), "FunctionType")),
278+
_ => None,
279+
}
280+
}
281+
282+
/// Create an empty `arguments` AST node (no parameters).
283+
fn empty_arguments_object(vm: &VirtualMachine) -> PyObjectRef {
284+
let node = NodeAst
285+
.into_ref_with_type(vm, pyast::NodeArguments::static_type().to_owned())
286+
.unwrap();
287+
let dict = node.as_object().dict().unwrap();
288+
for list_field in [
289+
"posonlyargs",
290+
"args",
291+
"kwonlyargs",
292+
"kw_defaults",
293+
"defaults",
294+
] {
295+
dict.set_item(list_field, vm.ctx.new_list(vec![]).into(), vm)
296+
.unwrap();
297+
}
298+
for none_field in ["vararg", "kwarg"] {
299+
dict.set_item(none_field, vm.ctx.none(), vm).unwrap();
300+
}
301+
node.into()
302+
}
303+
268304
#[cfg(feature = "parser")]
269305
pub(crate) fn parse(
270306
vm: &VirtualMachine,

‎crates/vm/src/stdlib/ast/expression.rs‎

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -330,31 +330,7 @@ impl Node for ast::ExprLambda {
330330
// Lambda with no parameters should have an empty arguments object, not None
331331
let args = match parameters {
332332
Some(params) => params.ast_to_object(vm, source_file),
333-
None => {
334-
// Create an empty arguments object
335-
let args_node = NodeAst
336-
.into_ref_with_type(vm, pyast::NodeArguments::static_type().to_owned())
337-
.unwrap();
338-
let args_dict = args_node.as_object().dict().unwrap();
339-
args_dict
340-
.set_item("posonlyargs", vm.ctx.new_list(vec![]).into(), vm)
341-
.unwrap();
342-
args_dict
343-
.set_item("args", vm.ctx.new_list(vec![]).into(), vm)
344-
.unwrap();
345-
args_dict.set_item("vararg", vm.ctx.none(), vm).unwrap();
346-
args_dict
347-
.set_item("kwonlyargs", vm.ctx.new_list(vec![]).into(), vm)
348-
.unwrap();
349-
args_dict
350-
.set_item("kw_defaults", vm.ctx.new_list(vec![]).into(), vm)
351-
.unwrap();
352-
args_dict.set_item("kwarg", vm.ctx.none(), vm).unwrap();
353-
args_dict
354-
.set_item("defaults", vm.ctx.new_list(vec![]).into(), vm)
355-
.unwrap();
356-
args_node.into()
357-
}
333+
None => empty_arguments_object(vm),
358334
};
359335
dict.set_item("args", args, vm).unwrap();
360336
dict.set_item("body", body.ast_to_object(vm, source_file), vm)

‎crates/vm/src/stdlib/ast/python.rs‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -95,12 +95,11 @@ pub(crate) mod _ast {
9595
.is_some();
9696
if has_field_types {
9797
// ASDL list fields (type*) default to empty list,
98-
// optional fields (type?) default to None.
98+
// optional/required fields default to None.
99+
// Fields that are always list-typed regardless of node class.
99100
const LIST_FIELDS: &[&str] = &[
100-
"args",
101101
"argtypes",
102102
"bases",
103-
"body",
104103
"cases",
105104
"comparators",
106105
"decorator_list",
@@ -113,31 +112,49 @@ pub(crate) mod _ast {
113112
"items",
114113
"keys",
115114
"kw_defaults",
115+
"kwd_attrs",
116+
"kwd_patterns",
116117
"keywords",
117118
"kwonlyargs",
118119
"names",
119-
"orelse",
120120
"ops",
121+
"patterns",
121122
"posonlyargs",
122123
"targets",
123124
"type_ignores",
124125
"type_params",
125126
"values",
126127
];
127128

129+
let class_name = zelf.class().name().to_string();
130+
128131
for field in &fields {
129132
if !set_fields.contains(field.as_str()) {
130-
let default: PyObjectRef = if LIST_FIELDS.contains(&field.as_str()) {
133+
let field_name = field.as_str();
134+
// Some field names have different ASDL types depending on the node.
135+
// For example, "args" is `expr*` in Call but `arguments` in Lambda.
136+
// "body" and "orelse" are `stmt*` in most nodes but `expr` in IfExp.
137+
let is_list_field = if field_name == "args" {
138+
class_name == "Call" || class_name == "arguments"
139+
} else if field_name == "body" || field_name == "orelse" {
140+
!matches!(
141+
class_name.as_str(),
142+
"Lambda" | "Expression" | "IfExp"
143+
)
144+
} else {
145+
LIST_FIELDS.contains(&field_name)
146+
};
147+
148+
let default: PyObjectRef = if is_list_field {
131149
vm.ctx.new_list(vec![]).into()
132150
} else {
133151
vm.ctx.none()
134152
};
135-
zelf.set_attr(vm.ctx.intern_str(field.as_str()), default, vm)?;
153+
zelf.set_attr(vm.ctx.intern_str(field_name), default, vm)?;
136154
}
137155
}
138156

139157
// Special defaults that are not None or empty list
140-
let class_name = &*zelf.class().name();
141158
if class_name == "ImportFrom" && !set_fields.contains("level") {
142159
zelf.set_attr("level", vm.ctx.new_int(0), vm)?;
143160
}

‎crates/vm/src/stdlib/builtins.rs‎

Lines changed: 24 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -142,24 +142,32 @@ mod builtins {
142142
.fast_isinstance(&ast::NodeAst::make_class(&vm.ctx))
143143
{
144144
use num_traits::Zero;
145-
let flags = args.flags.map_or(Ok(0), |v| v.try_to_primitive(vm))?;
145+
let flags: i32 = args.flags.map_or(Ok(0), |v| v.try_to_primitive(vm))?;
146+
let is_ast_only = !(flags & ast::PY_COMPILE_FLAG_AST_ONLY).is_zero();
147+
148+
// func_type mode requires PyCF_ONLY_AST
149+
if mode_str == "func_type" && !is_ast_only {
150+
return Err(vm.new_value_error(
151+
"compile() mode 'func_type' requires flag PyCF_ONLY_AST".to_owned(),
152+
));
153+
}
154+
146155
// compile(ast_node, ..., PyCF_ONLY_AST) returns the AST after validation
147-
if !(flags & ast::PY_COMPILE_FLAG_AST_ONLY).is_zero() {
148-
let expected_type = match mode_str {
149-
"exec" => "Module",
150-
"eval" => "Expression",
151-
"single" => "Interactive",
152-
"func_type" => "FunctionType",
153-
_ => {
154-
return Err(vm.new_value_error(format!(
155-
"compile() mode must be 'exec', 'eval', 'single' or 'func_type', got '{mode_str}'"
156-
)));
157-
}
158-
};
159-
let cls_name = args.source.class().name().to_string();
160-
if cls_name != expected_type {
156+
if is_ast_only {
157+
let (expected_type, expected_name) = ast::mode_type_and_name(
158+
&vm.ctx, mode_str,
159+
)
160+
.ok_or_else(|| {
161+
vm.new_value_error(
162+
"compile() mode must be 'exec', 'eval', 'single' or 'func_type'"
163+
.to_owned(),
164+
)
165+
})?;
166+
if !args.source.fast_isinstance(&expected_type) {
161167
return Err(vm.new_type_error(format!(
162-
"expected {expected_type} node, got {cls_name}"
168+
"expected {} node, got {}",
169+
expected_name,
170+
args.source.class().name()
163171
)));
164172
}
165173
return Ok(args.source);

0 commit comments

Comments
 (0)