diff --git a/crates/vm/src/stdlib/ast/parameter.rs b/crates/vm/src/stdlib/ast/parameter.rs index 1e411d41a..b1942d833 100644 --- a/crates/vm/src/stdlib/ast/parameter.rs +++ b/crates/vm/src/stdlib/ast/parameter.rs @@ -126,8 +126,9 @@ impl Node for ast::Parameter { _vm, ) .unwrap(); - // dict.set_item("type_comment", type_comment.ast_to_object(_vm), _vm) - // .unwrap(); + // Ruff AST doesn't track type_comment, so always set to None + dict.set_item("type_comment", _vm.ctx.none(), _vm) + .unwrap(); node_add_location(&dict, range, _vm, source_file); node.into() } diff --git a/crates/vm/src/stdlib/ast/pyast.rs b/crates/vm/src/stdlib/ast/pyast.rs index 2131df29b..d6f995f6f 100644 --- a/crates/vm/src/stdlib/ast/pyast.rs +++ b/crates/vm/src/stdlib/ast/pyast.rs @@ -37,6 +37,12 @@ macro_rules! impl_node { ),* ]).into(), ); + + // Signal that this is a built-in AST node with field defaults + class.set_attr( + ctx.intern_str("_field_types"), + ctx.new_dict().into(), + ); } } }; @@ -902,21 +908,21 @@ impl_node!( impl_node!( #[pyclass(module = "_ast", name = "TypeVar", base = NodeTypeParam)] pub(crate) struct NodeTypeParamTypeVar, - fields: ["name", "bound"], + fields: ["name", "bound", "default_value"], attributes: ["lineno", "col_offset", "end_lineno", "end_col_offset"], ); impl_node!( #[pyclass(module = "_ast", name = "ParamSpec", base = NodeTypeParam)] pub(crate) struct NodeTypeParamParamSpec, - fields: ["name"], + fields: ["name", "default_value"], attributes: ["lineno", "col_offset", "end_lineno", "end_col_offset"], ); impl_node!( #[pyclass(module = "_ast", name = "TypeVarTuple", base = NodeTypeParam)] pub(crate) struct NodeTypeParamTypeVarTuple, - fields: ["name"], + fields: ["name", "default_value"], attributes: ["lineno", "col_offset", "end_lineno", "end_col_offset"], ); diff --git a/crates/vm/src/stdlib/ast/python.rs b/crates/vm/src/stdlib/ast/python.rs index 152ae9abc..924026735 100644 --- a/crates/vm/src/stdlib/ast/python.rs +++ b/crates/vm/src/stdlib/ast/python.rs @@ -87,24 +87,56 @@ pub(crate) mod _ast { zelf.set_attr(vm.ctx.intern_str(key), value, vm)?; } - // Set default values for fields that weren't provided - let class_name = &*zelf.class().name(); - if class_name == "Module" && !set_fields.contains("type_ignores") { - zelf.set_attr("type_ignores", vm.ctx.new_list(vec![]), vm)?; - } - if class_name == "ImportFrom" && !set_fields.contains("level") { - zelf.set_attr("level", vm.ctx.new_int(0), vm)?; - } - if class_name == "alias" && !set_fields.contains("asname") { - zelf.set_attr("asname", vm.ctx.none(), vm)?; - } - // Set type_comment to None for nodes that support it - if !set_fields.contains("type_comment") { - match class_name { - "FunctionDef" | "AsyncFunctionDef" | "For" | "AsyncFor" | "With" | "AsyncWith" | "arg" => { - zelf.set_attr("type_comment", vm.ctx.none(), vm)?; + // Set default values only for built-in AST nodes (_field_types present). + // Custom AST subclasses without _field_types do NOT get automatic defaults. + let has_field_types = zelf.class().get_attr(vm.ctx.intern_str("_field_types")).is_some(); + if has_field_types { + // ASDL list fields (type*) default to empty list, + // optional fields (type?) default to None. + const LIST_FIELDS: &[&str] = &[ + "args", + "argtypes", + "bases", + "body", + "cases", + "comparators", + "decorator_list", + "defaults", + "elts", + "finalbody", + "generators", + "handlers", + "ifs", + "items", + "keys", + "kw_defaults", + "keywords", + "kwonlyargs", + "names", + "orelse", + "ops", + "posonlyargs", + "targets", + "type_ignores", + "type_params", + "values", + ]; + + for field in &fields { + if !set_fields.contains(field.as_str()) { + let default: PyObjectRef = if LIST_FIELDS.contains(&field.as_str()) { + vm.ctx.new_list(vec![]).into() + } else { + vm.ctx.none() + }; + zelf.set_attr(vm.ctx.intern_str(field.as_str()), default, vm)?; } - _ => {} + } + + // Special defaults that are not None or empty list + let class_name = &*zelf.class().name(); + if class_name == "ImportFrom" && !set_fields.contains("level") { + zelf.set_attr("level", vm.ctx.new_int(0), vm)?; } } diff --git a/crates/vm/src/stdlib/ast/statement.rs b/crates/vm/src/stdlib/ast/statement.rs index b7bc692dd..620a83178 100644 --- a/crates/vm/src/stdlib/ast/statement.rs +++ b/crates/vm/src/stdlib/ast/statement.rs @@ -182,9 +182,9 @@ impl Node for ast::StmtFunctionDef { .unwrap(); dict.set_item("returns", returns.ast_to_object(vm, source_file), vm) .unwrap(); - // TODO: Ruff ignores type_comment during parsing - // dict.set_item("type_comment", type_comment.ast_to_object(_vm), _vm) - // .unwrap(); + // Ruff AST doesn't track type_comment, so always set to None + dict.set_item("type_comment", vm.ctx.none(), vm) + .unwrap(); dict.set_item( "type_params", type_params @@ -647,8 +647,9 @@ impl Node for ast::StmtFor { .unwrap(); dict.set_item("orelse", orelse.ast_to_object(_vm, source_file), _vm) .unwrap(); - // dict.set_item("type_comment", type_comment.ast_to_object(_vm), _vm) - // .unwrap(); + // Ruff AST doesn't track type_comment, so always set to None + dict.set_item("type_comment", _vm.ctx.none(), _vm) + .unwrap(); node_add_location(&dict, _range, _vm, source_file); node.into() } @@ -799,8 +800,9 @@ impl Node for ast::StmtWith { .unwrap(); dict.set_item("body", body.ast_to_object(_vm, source_file), _vm) .unwrap(); - // dict.set_item("type_comment", type_comment.ast_to_object(_vm), _vm) - // .unwrap(); + // Ruff AST doesn't track type_comment, so always set to None + dict.set_item("type_comment", _vm.ctx.none(), _vm) + .unwrap(); node_add_location(&dict, _range, _vm, source_file); node.into() } diff --git a/crates/vm/src/stdlib/ast/type_parameters.rs b/crates/vm/src/stdlib/ast/type_parameters.rs index 4801a9a4b..ccfbf4649 100644 --- a/crates/vm/src/stdlib/ast/type_parameters.rs +++ b/crates/vm/src/stdlib/ast/type_parameters.rs @@ -78,7 +78,7 @@ impl Node for ast::TypeParamTypeVar { name, bound, range: _range, - default: _, + default, } = self; let node = NodeAst .into_ref_with_type(_vm, pyast::NodeTypeParamTypeVar::static_type().to_owned()) @@ -88,6 +88,12 @@ impl Node for ast::TypeParamTypeVar { .unwrap(); dict.set_item("bound", bound.ast_to_object(_vm, source_file), _vm) .unwrap(); + dict.set_item( + "default_value", + default.ast_to_object(_vm, source_file), + _vm, + ) + .unwrap(); node_add_location(&dict, _range, _vm, source_file); node.into() } diff --git a/crates/vm/src/stdlib/builtins.rs b/crates/vm/src/stdlib/builtins.rs index 6c72c4c69..8bfbffcc6 100644 --- a/crates/vm/src/stdlib/builtins.rs +++ b/crates/vm/src/stdlib/builtins.rs @@ -141,10 +141,27 @@ mod builtins { .source .fast_isinstance(&ast::NodeAst::make_class(&vm.ctx)) { - // If PyCF_ONLY_AST is set, just return the AST node as-is use num_traits::Zero; let flags = args.flags.map_or(Ok(0), |v| v.try_to_primitive(vm))?; + // compile(ast_node, ..., PyCF_ONLY_AST) returns the AST after validation if !(flags & ast::PY_COMPILE_FLAG_AST_ONLY).is_zero() { + let expected_type = match mode_str { + "exec" => "Module", + "eval" => "Expression", + "single" => "Interactive", + "func_type" => "FunctionType", + _ => { + return Err(vm.new_value_error(format!( + "compile() mode must be 'exec', 'eval', 'single' or 'func_type', got '{mode_str}'" + ))); + } + }; + let cls_name = args.source.class().name().to_string(); + if cls_name != expected_type { + return Err(vm.new_type_error(format!( + "expected {expected_type} node, got {cls_name}" + ))); + } return Ok(args.source); }