Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
28 changes: 10 additions & 18 deletions crates/vm/src/stdlib/_ast/pyast.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,17 +72,13 @@ macro_rules! impl_base_node {
super::python::_ast::ast_reduce(&zelf, vm)
}

#[pymethod]
fn __replace__(
zelf: PyObjectRef,
fields: crate::function::KwArgs<PyObjectRef, crate::function::NameFields>,
vm: &VirtualMachine,
) -> PyResult {
super::python::_ast::ast_replace(&zelf, fields.into(), vm)
}

#[extend_class]
fn extend_class(ctx: &Context, class: &'static Py<PyType>) {
class.set_str_attr(
"__replace__",
super::python::_ast::AST_REPLACE.to_proper_method(class, ctx),
ctx,
);
// AST types are mutable (heap types, not IMMUTABLETYPE).
class.slots.flags.remove(crate::types::PyTypeFlags::IMMUTABLETYPE);
class.set_attr(
Expand Down Expand Up @@ -116,17 +112,13 @@ macro_rules! impl_base_node {
super::python::_ast::ast_reduce(&zelf, vm)
}

#[pymethod]
fn __replace__(
zelf: PyObjectRef,
fields: crate::function::KwArgs<PyObjectRef, crate::function::NameFields>,
vm: &VirtualMachine,
) -> PyResult {
super::python::_ast::ast_replace(&zelf, fields.into(), vm)
}

#[extend_class]
fn extend_class_with_fields(ctx: &Context, class: &'static Py<PyType>) {
class.set_str_attr(
"__replace__",
super::python::_ast::AST_REPLACE.to_proper_method(class, ctx),
ctx,
);
// AST types are mutable (heap types, not IMMUTABLETYPE).
class.slots.flags.remove(crate::types::PyTypeFlags::IMMUTABLETYPE);
class.set_attr(
Expand Down
129 changes: 71 additions & 58 deletions crates/vm/src/stdlib/_ast/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ pub(crate) mod _ast {
types::{Constructor, Initializer},
warn,
};
use rustpython_common::wtf8::Wtf8Buf;

#[pyattr]
#[pyclass(module = "ast", name = "AST")]
#[derive(Debug, PyPayload)]
Expand Down Expand Up @@ -167,6 +169,17 @@ pub(crate) mod _ast {
obj.set_attr(&name, value, vm)
}

pub(crate) const AST_REPLACE: PyMethodDef = PyMethodDef::new_const(
"__replace__",
|zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine| -> PyResult {
ast_replace(&zelf, args, vm)
},
PyMethodFlags::METHOD,
crate::function::ItemDoc::static_text(
"__replace__($self, /, **fields)\n--\n\nReturn a copy of the AST node with new values for the specified fields.",
),
);

pub(crate) fn ast_replace(zelf: &PyObject, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
if !args.args.is_empty() {
return Err(vm.new_type_error("__replace__() takes no positional arguments"));
Expand Down Expand Up @@ -413,68 +426,78 @@ This will become an error in Python 3.20",
zelf.set_attr(vm.ctx.intern_str(key), value, vm)?;
}

// Use _field_types to determine defaults for unset fields.
// Only built-in AST node classes have _field_types populated.
let field_types = zelf.class().get_attr(vm.ctx.intern_str("_field_types"));
if let Some(Ok(ft_dict)) =
field_types.map(|ft| ft.downcast::<crate::builtins::PyDict>())
// Custom AST subclasses without _field_types retain their permissive
// treatment of fields omitted from the constructor.
if !remaining_fields.elements().is_empty()
&& let Some(field_types) = zelf.class().get_attr(vm.ctx.intern_str("_field_types"))
{
let ft_dict = field_types
.downcast::<PyDict>()
.map_err(|_| vm.new_system_error("bad argument to internal function"))?;
let expr_ctx_type: PyObjectRef =
super::super::pyast::NodeExprContext::make_static_type().into();
let missing_names = PySet::default().into_ref(&vm.ctx);

let mut missing = Vec::new();
// Render missing arguments in their declared positional order.
for i in 0..numfields {
let field = fields_seq.get_item(i as isize, vm)?;
if !ast_replace_set_discard(&remaining_fields, &field, vm)? {
continue;
}
if let Some(ftype) = ft_dict.get_item_opt(&*field, vm)? {
if ftype.fast_isinstance(vm.ctx.types.union_type) {
// Optional field (T | None) — no default
} else if ftype.fast_isinstance(vm.ctx.types.generic_alias_type) {
// List field (list[T]) — default to []
let empty_list: PyObjectRef = vm.ctx.new_list(vec![]).into();
ast_set_attr(zelf, &field, empty_list, vm)?;
} else if ftype.is(&expr_ctx_type) {
// expr_context — default to Load()
let load_type =
super::super::pyast::NodeExprContextLoad::make_static_type();
let load_instance = load_type
.get_attr(vm.ctx.intern_str("_instance"))
.unwrap_or_else(|| {
vm.ctx.new_base_object(load_type, Some(vm.ctx.new_dict()))
});
ast_set_attr(zelf, &field, load_instance, vm)?;
} else {
missing.push(format!("'{}'", field.str(vm)?));
}
} else {
for field in remaining_fields.elements() {
let Some(ftype) = ft_dict.get_item_opt(&*field, vm)? else {
return Err(vm.new_type_error(format!(
"Field {} is missing from {}._field_types",
field.repr(vm)?,
zelf.class().fully_qualified_name(vm)?
)));
};
if ftype.fast_isinstance(vm.ctx.types.union_type) {
// Optional fields already have a None default on the class.
} else if ftype.class().is(vm.ctx.types.generic_alias_type) {
let empty_list: PyObjectRef = vm.ctx.new_list(vec![]).into();
ast_set_attr(zelf, &field, empty_list, vm)?;
} else if ftype.is(&expr_ctx_type) {
let load_type =
super::super::pyast::NodeExprContextLoad::make_static_type();
let load_instance = load_type
.get_attr(vm.ctx.intern_str("_instance"))
.unwrap_or_else(|| {
vm.ctx.new_base_object(load_type, Some(vm.ctx.new_dict()))
});
ast_set_attr(zelf, &field, load_instance, vm)?;
} else {
missing_names.add(field, vm)?;
}
}
if !missing.is_empty() {
let count = missing.len();
let names = match missing.as_slice() {
[name] => name.clone(),
[first, last] => format!("{first} and {last}"),
names => format!(
"{}, and {}",
names[..count - 1].join(", "),
names[count - 1]
),
};
return Err(vm.new_type_error(format!(
"{}.__init__ missing {} required positional argument{}: {}",

let count = missing_names.elements().len();
if count > 0 {
// Match constructor order and preserve non-UTF-8 field names.
let mut message = Wtf8Buf::from(format!(
"{}.__init__ missing {} required positional argument{}: ",
zelf.class().fully_qualified_name(vm)?,
count,
if count == 1 { "" } else { "s" },
names
)));
));
let mut num_left = count as isize;
for i in 0..numfields {
let field = fields_seq.get_item(i as isize, vm)?;
if missing_names.contains(&field, vm)? {
let field = field
.downcast::<PyStr>()
.map_err(|_| vm.new_system_error("non-string AST field name"))?;
message.push_wtf8("'".as_ref());
message.push_wtf8(field.as_wtf8());
message.push_wtf8("'".as_ref());
let separator = if num_left == 1 {
""
} else if count == 2 {
" and "
} else if num_left == 2 {
", and "
} else {
", "
};
message.push_wtf8(separator.as_ref());
num_left -= 1;
}
}
return Err(vm.new_type_error(message));
}
}

Expand Down Expand Up @@ -555,16 +578,6 @@ This will become an error in Python 3.20",
PyMethodFlags::METHOD,
crate::function::ItemDoc::static_text("__reduce__($self, /)\n--\n\n"),
);
const AST_REPLACE: PyMethodDef = PyMethodDef::new_const(
"__replace__",
|zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine| -> PyResult {
ast_replace(&zelf, args, vm)
},
PyMethodFlags::METHOD,
crate::function::ItemDoc::static_text(
"__replace__($self, /, **fields)\n--\n\nReturn a copy of the AST node with new values for the specified fields.",
),
);
let base_type = NodeAst::static_type();
ast_type.set_str_attr(
"__reduce__",
Expand Down
Loading