diff --git a/crates/vm/src/exceptions.rs b/crates/vm/src/exceptions.rs index f7e1b79aa5a..1db10274e89 100644 --- a/crates/vm/src/exceptions.rs +++ b/crates/vm/src/exceptions.rs @@ -959,9 +959,7 @@ impl ExceptionZoo { "exceptions" => ctx.new_readonly_getset("exceptions", excs.base_exception_group, make_arg_getter(1)), }); - extend_exception!(PySystemExit, ctx, excs.system_exit, { - "code" => ctx.new_readonly_getset("code", excs.system_exit, system_exit_code), - }); + extend_exception!(PySystemExit, ctx, excs.system_exit); extend_exception!(PyKeyboardInterrupt, ctx, excs.keyboard_interrupt); extend_exception!(PyGeneratorExit, ctx, excs.generator_exit); @@ -1100,19 +1098,6 @@ fn syntax_error_set_msg(exc: PyBaseExceptionRef, value: PySetterValue, vm: &Virt *args = PyTuple::new_ref(new_args, &vm.ctx); } -fn system_exit_code(exc: PyBaseExceptionRef) -> Option { - // SystemExit.code based on args length: - // - size == 0: code is None - // - size == 1: code is args[0] - // - size > 1: code is args (the whole tuple) - let args = exc.args.read(); - Some(match args.len() { - 0 => return None, - 1 => args.first().unwrap().clone(), - _ => args.as_object().to_owned(), - }) -} - #[cfg(feature = "serde")] pub struct SerializeException<'vm, 's> { vm: &'vm VirtualMachine, @@ -1593,7 +1578,7 @@ pub(super) mod types { tuple::IntoPyTuple, }, convert::ToPyResult, - function::{ArgBytesLike, FuncArgs, KwArgs}, + function::{ArgBytesLike, FuncArgs, KwArgs, PySetterValue}, set_attrs, types::{Constructor, Initializer}, }; @@ -1624,22 +1609,65 @@ pub(super) mod types { pub(super) args: PyRwLock, } - #[pyexception(name, base = PyBaseException, ctx = "system_exit")] - #[derive(Debug)] - #[repr(transparent)] - pub struct PySystemExit(PyBaseException); + #[pyexception(name, base = PyBaseException, ctx = "system_exit", traverse = "manual")] + #[repr(C)] + pub struct PySystemExit { + base: PyBaseException, + code: PyAtomicRef>, + } - // SystemExit_init: has its own __init__ that sets the code attribute - #[pyexception(with(Initializer))] - impl PySystemExit {} + impl crate::class::PySubclass for PySystemExit { + type Base = PyBaseException; + fn as_base(&self) -> &Self::Base { + &self.base + } + } + + unsafe impl Traverse for PySystemExit { + fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) { + self.base.traverse(tracer_fn); + if let Some(obj) = self.code.deref() { + tracer_fn(obj); + } + } + } + + impl core::fmt::Debug for PySystemExit { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("PySystemExit").finish_non_exhaustive() + } + } + + #[pyexception(with(Constructor, Initializer))] + impl PySystemExit { + #[pygetset] + fn code(&self) -> Option { + self.code.to_owned() + } + + #[pygetset(setter)] + fn set_code(&self, value: PySetterValue, vm: &VirtualMachine) { + let code = match value { + PySetterValue::Assign(v) => Some(v), + PySetterValue::Delete => None, + }; + self.code.swap_to_temporary_refs(code, vm); + } + } impl Initializer for PySystemExit { type Args = FuncArgs; fn slot_init(zelf: PyObjectRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> { // Call BaseException_init first (handles args) - PyBaseException::slot_init(zelf, args, vm) - // Note: code is computed dynamically via system_exit_code getter - // so we don't need to set it here explicitly + let code = match args.args.len() { + 0 => vm.ctx.none(), + 1 => args.args[0].clone(), + _ => vm.ctx.new_tuple(args.args.clone()).into(), + }; + PyBaseException::slot_init(zelf.clone(), args, vm)?; + let exc: &Py = zelf.downcast_ref::().unwrap(); + exc.code.swap_to_temporary_refs(Some(code), vm); + Ok(()) } fn init(_zelf: PyRef, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<()> { @@ -1647,6 +1675,18 @@ pub(super) mod types { } } + impl Constructor for PySystemExit { + type Args = FuncArgs; + + fn py_new(_cls: &Py, args: FuncArgs, vm: &VirtualMachine) -> PyResult { + let base_exception = PyBaseException::new(args.args, vm); + Ok(Self { + base: base_exception, + code: None.into(), + }) + } + } + #[pyexception(name, base = PyBaseException, ctx = "generator_exit", impl)] #[derive(Debug)] #[repr(transparent)] diff --git a/crates/vm/src/stdlib/_thread.rs b/crates/vm/src/stdlib/_thread.rs index 99f9b1787ed..9caa15dbfee 100644 --- a/crates/vm/src/stdlib/_thread.rs +++ b/crates/vm/src/stdlib/_thread.rs @@ -623,7 +623,7 @@ pub(crate) mod _thread { #[pyfunction] fn exit(vm: &VirtualMachine) -> PyResult { - Err(vm.new_exception_empty(vm.ctx.exceptions.system_exit.to_owned())) + Err(vm.invoke_exception(vm.ctx.exceptions.system_exit.to_owned(), vec![])?) } thread_local!(static SENTINELS: RefCell>> = const { RefCell::new(Vec::new()) }); diff --git a/crates/vm/src/stdlib/builtins.rs b/crates/vm/src/stdlib/builtins.rs index 8ea849ab05b..27f30158b22 100644 --- a/crates/vm/src/stdlib/builtins.rs +++ b/crates/vm/src/stdlib/builtins.rs @@ -1041,7 +1041,7 @@ mod builtins { #[pyfunction] pub(super) fn exit(exit_code_arg: OptionalArg, vm: &VirtualMachine) -> PyResult { let code = exit_code_arg.unwrap_or_else(|| vm.ctx.new_int(0).into()); - Err(vm.new_exception(vm.ctx.exceptions.system_exit.to_owned(), vec![code])) + Err(vm.invoke_exception(vm.ctx.exceptions.system_exit.to_owned(), vec![code])?) } #[derive(Debug, Default, FromArgs)] diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 852e57cc0b4..0f2e6a46370 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -2161,7 +2161,7 @@ impl VirtualMachine { if self.state.finalizing.load(Ordering::Acquire) && !self.is_main_thread() { // once finalization starts, // non-main Python threads should stop running bytecode. - return Err(self.new_exception(self.ctx.exceptions.system_exit.to_owned(), vec![])); + return Err(self.invoke_exception(self.ctx.exceptions.system_exit.to_owned(), vec![])?); } // Suspend this thread if stop-the-world is in progress @@ -2303,29 +2303,28 @@ impl VirtualMachine { pub fn handle_exit_exception(&self, exc: PyBaseExceptionRef) -> u32 { if exc.fast_isinstance(self.ctx.exceptions.system_exit) { - let args = exc.args(); - let msg = match args.as_slice() { - [] => return 0, - [arg] => match_class!(match arg { - ref i @ PyInt => { - use num_traits::cast::ToPrimitive; - // Try u32 first, then i32 (for negative values), else -1 for overflow - let code = i - .as_bigint() - .to_u32() - .or_else(|| i.as_bigint().to_i32().map(|v| v as u32)) - .unwrap_or(-1i32 as u32); - return code; - } - arg => { - if self.is_none(arg) { - return 0; - } - arg.str(self).ok() + let code = exc + .as_object() + .get_attr("code", self) + .unwrap_or_else(|_| exc.as_object().to_owned()); + let msg = match_class!(match code { + ref i @ PyInt => { + use num_traits::cast::ToPrimitive; + // Try u32 first, then i32 (for negative values), else -1 for overflow + let code = i + .as_bigint() + .to_u32() + .or_else(|| i.as_bigint().to_i32().map(|v| v as u32)) + .unwrap_or(-1i32 as u32); + return code; + } + code => { + if self.is_none(&code) { + return 0; } - }), - _ => args.as_object().repr(self).ok(), - }; + code.str(self).ok() + } + }); if let Some(msg) = msg { // Write using Python's write() to use stderr's error handler (backslashreplace) if let Ok(stderr) = stdlib::sys::get_stderr(self) {