diff --git a/crates/vm/src/exceptions.rs b/crates/vm/src/exceptions.rs index b91e567e731..cf83b416f99 100644 --- a/crates/vm/src/exceptions.rs +++ b/crates/vm/src/exceptions.rs @@ -1593,6 +1593,7 @@ pub(super) mod types { }, convert::ToPyResult, function::{ArgBytesLike, FuncArgs, KwArgs}, + set_attrs, types::{Constructor, Initializer}, }; use crossbeam_utils::atomic::AtomicCell; @@ -1694,6 +1695,7 @@ pub(super) mod types { #[derive(Debug)] #[repr(transparent)] pub struct PyFloatingPointError(PyArithmeticError); + #[pyexception(name, base = PyArithmeticError, ctx = "overflow_error", impl)] #[derive(Debug)] #[repr(transparent)] @@ -1738,8 +1740,10 @@ pub(super) mod types { PyBaseException::slot_init(zelf.clone(), base_args, vm)?; // Set attributes - zelf.set_attr("name", vm.unwrap_or_none(name), vm)?; - zelf.set_attr("obj", vm.unwrap_or_none(obj), vm)?; + set_attrs!(zelf, vm, + "name" => vm.unwrap_or_none(name), + "obj" => vm.unwrap_or_none(obj), + ); Ok(()) } @@ -2653,12 +2657,13 @@ pub(super) mod types { fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> { type Args = (PyStrRef, ArgBytesLike, isize, isize, PyStrRef); let (encoding, object, start, end, reason): Args = args.bind(vm)?; - zelf.set_attr("encoding", encoding, vm)?; - let object_as_bytes = vm.ctx.new_bytes(object.borrow_buf().to_vec()); - zelf.set_attr("object", object_as_bytes, vm)?; - zelf.set_attr("start", vm.ctx.new_int(start), vm)?; - zelf.set_attr("end", vm.ctx.new_int(end), vm)?; - zelf.set_attr("reason", reason, vm)?; + set_attrs!(zelf, vm, + "encoding" => encoding, + "object" => vm.ctx.new_bytes(object.borrow_buf().to_vec()), + "start" => vm.ctx.new_int(start), + "end" => vm.ctx.new_int(end), + "reason" => reason, + ); Ok(()) } @@ -2711,11 +2716,13 @@ pub(super) mod types { fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> { type Args = (PyStrRef, PyStrRef, isize, isize, PyStrRef); let (encoding, object, start, end, reason): Args = args.bind(vm)?; - zelf.set_attr("encoding", encoding, vm)?; - zelf.set_attr("object", object, vm)?; - zelf.set_attr("start", vm.ctx.new_int(start), vm)?; - zelf.set_attr("end", vm.ctx.new_int(end), vm)?; - zelf.set_attr("reason", reason, vm)?; + set_attrs!(zelf, vm, + "encoding" => encoding, + "object" => object, + "start" => vm.ctx.new_int(start), + "end" => vm.ctx.new_int(end), + "reason" => reason, + ); Ok(()) } @@ -2766,10 +2773,12 @@ pub(super) mod types { fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> { type Args = (PyStrRef, isize, isize, PyStrRef); let (object, start, end, reason): Args = args.bind(vm)?; - zelf.set_attr("object", object, vm)?; - zelf.set_attr("start", vm.ctx.new_int(start), vm)?; - zelf.set_attr("end", vm.ctx.new_int(end), vm)?; - zelf.set_attr("reason", reason, vm)?; + set_attrs!(zelf, vm, + "object" => object, + "start" => vm.ctx.new_int(start), + "end" => vm.ctx.new_int(end), + "reason" => reason, + ); Ok(()) } diff --git a/crates/vm/src/protocol/object.rs b/crates/vm/src/protocol/object.rs index 0d4e206ec19..3398a48e7aa 100644 --- a/crates/vm/src/protocol/object.rs +++ b/crates/vm/src/protocol/object.rs @@ -826,3 +826,18 @@ impl PyObject { } } } + +/// Macro to reduce code repetition when setting multiple attributes on an object. +#[macro_export] +macro_rules! set_attrs { + ($obj:expr, $vm:expr, unwrap, $($key:expr => $val:expr),+ $(,)?) => { + $( + $obj.set_attr($key, $val, $vm).unwrap(); + )+ + }; + ($obj:expr, $vm:expr, $($key:expr => $val:expr),+ $(,)?) => { + $( + $obj.set_attr($key, $val, $vm)?; + )+ + }; +} diff --git a/crates/vm/src/stdlib/_ctypes/simple.rs b/crates/vm/src/stdlib/_ctypes/simple.rs index d51a61130f3..122c23cc25c 100644 --- a/crates/vm/src/stdlib/_ctypes/simple.rs +++ b/crates/vm/src/stdlib/_ctypes/simple.rs @@ -11,7 +11,9 @@ use crate::convert::ToPyObject; use crate::function::{Either, FuncArgs, OptionalArg}; use crate::protocol::{BufferDescriptor, PyBuffer, PyNumberMethods}; use crate::types::{AsBuffer, AsNumber, Constructor, Initializer, Representable}; -use crate::{AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine}; +use crate::{ + AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, set_attrs, +}; use alloc::borrow::Cow; use core::fmt::Debug; use num_traits::ToPrimitive; @@ -605,12 +607,11 @@ fn create_swapped_types( // c_char (c), c_byte (b), c_ubyte (B) let single_byte_types = ["c", "b", "B"]; if single_byte_types.contains(&type_str) { - type_ref - .as_object() - .set_attr("__ctype_le__", type_ref.as_object().to_owned(), vm)?; - type_ref - .as_object() - .set_attr("__ctype_be__", type_ref.as_object().to_owned(), vm)?; + set_attrs!( + type_ref.as_object(), vm, + "__ctype_le__" => type_ref.as_object().to_owned(), + "__ctype_be__" => type_ref.as_object().to_owned(), + ); return Ok(()); } @@ -627,12 +628,11 @@ fn create_swapped_types( let bases = vm.ctx.new_tuple(vec![type_ref.as_object().to_owned()]); // Set placeholder first to prevent recursion - type_ref - .as_object() - .set_attr("__ctype_le__", vm.ctx.none(), vm)?; - type_ref - .as_object() - .set_attr("__ctype_be__", vm.ctx.none(), vm)?; + set_attrs!( + type_ref.as_object(), vm, + "__ctype_le__" => vm.ctx.none(), + "__ctype_be__" => vm.ctx.none(), + ); // Create only the non-native endian type let suffix = if is_little_endian { "_be" } else { "_le" }; @@ -646,7 +646,10 @@ fn create_swapped_types( )?; // Set _swappedbytes_ on the swapped type to indicate byte swapping is needed - swapped_type.set_attr("_swappedbytes_", vm.ctx.none(), vm)?; + set_attrs!( + swapped_type, vm, + "_swappedbytes_" => vm.ctx.none(), + ); // Update swapped type's StgInfo format to use opposite endian prefix if let Ok(swapped_type_ref) = swapped_type.clone().downcast::() @@ -661,27 +664,27 @@ fn create_swapped_types( // Set attributes based on system byte order // Native endian attribute points to self, non-native points to swapped type - if is_little_endian { + let type_obj = type_ref.as_object().to_owned(); + let swapped_obj = swapped_type.clone(); + let (ctype_le, ctype_be) = if is_little_endian { // Little-endian system: __ctype_le__ = self, __ctype_be__ = swapped - type_ref - .as_object() - .set_attr("__ctype_le__", type_ref.as_object().to_owned(), vm)?; - type_ref - .as_object() - .set_attr("__ctype_be__", swapped_type.clone(), vm)?; - swapped_type.set_attr("__ctype_le__", type_ref.as_object().to_owned(), vm)?; - swapped_type.set_attr("__ctype_be__", swapped_type.clone(), vm)?; + (type_obj, swapped_obj) } else { - // Big-endian system: __ctype_be__ = self, __ctype_le__ = swapped - type_ref - .as_object() - .set_attr("__ctype_be__", type_ref.as_object().to_owned(), vm)?; - type_ref - .as_object() - .set_attr("__ctype_le__", swapped_type.clone(), vm)?; - swapped_type.set_attr("__ctype_be__", type_ref.as_object().to_owned(), vm)?; - swapped_type.set_attr("__ctype_le__", swapped_type.clone(), vm)?; - } + // Big-endian system: __ctype_le__ = swapped, __ctype_be__ = self + (swapped_obj, type_obj) + }; + + set_attrs!( + type_ref.as_object(), vm, + "__ctype_le__" => ctype_le.clone(), + "__ctype_be__" => ctype_be.clone(), + ); + + set_attrs!( + swapped_type, vm, + "__ctype_le__" => ctype_le, + "__ctype_be__" => ctype_be, + ); Ok(()) } diff --git a/crates/vm/src/vm/vm_new.rs b/crates/vm/src/vm/vm_new.rs index 3b50c25695f..5af053faffc 100644 --- a/crates/vm/src/vm/vm_new.rs +++ b/crates/vm/src/vm/vm_new.rs @@ -22,6 +22,7 @@ use crate::{ exceptions::OSErrorBuilder, function::{IntoPyNativeFn, PyMethodFlags}, scope::Scope, + set_attrs, vm::VirtualMachine, }; @@ -425,7 +426,11 @@ impl VirtualMachine { pub fn new_name_error(&self, msg: impl Into, name: PyStrRef) -> PyBaseExceptionRef { let name_error_type = self.ctx.exceptions.name_error.to_owned(); let name_error = self.new_exception_msg(name_error_type, msg.into()); - name_error.as_object().set_attr("name", name, self).unwrap(); + set_attrs!( + name_error.as_object(), self, unwrap, + "name" => name, + ); + name_error } @@ -512,13 +517,16 @@ impl VirtualMachine { reason.clone().into(), ], ); - exc.as_object() - .set_attr("encoding", encoding, self) - .unwrap(); - exc.as_object().set_attr("object", object, self).unwrap(); - exc.as_object().set_attr("start", start, self).unwrap(); - exc.as_object().set_attr("end", end, self).unwrap(); - exc.as_object().set_attr("reason", reason, self).unwrap(); + + set_attrs!( + exc.as_object(), self, unwrap, + "encoding" => encoding, + "object" => object, + "start" => start, + "end" => end, + "reason" => reason, + ); + exc } @@ -542,13 +550,16 @@ impl VirtualMachine { reason.clone().into(), ], ); - exc.as_object() - .set_attr("encoding", encoding, self) - .unwrap(); - exc.as_object().set_attr("object", object, self).unwrap(); - exc.as_object().set_attr("start", start, self).unwrap(); - exc.as_object().set_attr("end", end, self).unwrap(); - exc.as_object().set_attr("reason", reason, self).unwrap(); + + set_attrs!( + exc.as_object(), self, unwrap, + "encoding" => encoding, + "object" => object, + "start" => start, + "end" => end, + "reason" => reason, + ); + exc } @@ -752,14 +763,12 @@ impl VirtualMachine { let (lineno, offset) = error.python_location(); let lineno = self.ctx.new_int(lineno); let offset = self.ctx.new_int(offset); - syntax_error - .as_object() - .set_attr("lineno", lineno, self) - .unwrap(); - syntax_error - .as_object() - .set_attr("offset", offset, self) - .unwrap(); + + set_attrs!( + syntax_error.as_object(), self, unwrap, + "lineno" => lineno, + "offset" => offset, + ); // Set end_lineno and end_offset if available if let Some((end_lineno, end_offset)) = error.python_end_location() { @@ -771,24 +780,19 @@ impl VirtualMachine { }; let end_lineno = self.ctx.new_int(end_lineno); let end_offset = self.ctx.new_int(end_offset); - syntax_error - .as_object() - .set_attr("end_lineno", end_lineno, self) - .unwrap(); - syntax_error - .as_object() - .set_attr("end_offset", end_offset, self) - .unwrap(); + + set_attrs!( + syntax_error.as_object(), self, unwrap, + "end_lineno" => end_lineno, + "end_offset" => end_offset, + ); } - syntax_error - .as_object() - .set_attr("text", statement.to_pyobject(self), self) - .unwrap(); - syntax_error - .as_object() - .set_attr("filename", self.ctx.new_str(error.source_path()), self) - .unwrap(); + set_attrs!( + syntax_error.as_object(), self, unwrap, + "text" => statement.to_pyobject(self), + "filename" => self.ctx.new_str(error.source_path()) + ); // Set _metadata for keyword typo suggestions in traceback module. // Format: (start_line, col_offset, source_code) @@ -800,10 +804,10 @@ impl VirtualMachine { self.ctx.new_int(0).into(), self.ctx.new_str(source).into(), ]); - syntax_error - .as_object() - .set_attr("_metadata", metadata, self) - .unwrap(); + set_attrs!( + syntax_error.as_object(), self, unwrap, + "_metadata" => metadata, + ); } syntax_error @@ -825,7 +829,11 @@ impl VirtualMachine { ) -> PyBaseExceptionRef { let import_error = self.ctx.exceptions.import_error.to_owned(); let exc = self.new_exception_msg(import_error, msg.into()); - exc.as_object().set_attr("name", name.into(), self).unwrap(); + set_attrs!( + exc.as_object(), self, unwrap, + "name" => name.into(), + ); + exc }