From ee2dff68b30a639bcf06b8de9cc2664cf62a1998 Mon Sep 17 00:00:00 2001 From: Noah <33094578+coolreader18@users.noreply.github.com> Date: Wed, 16 Sep 2020 09:26:22 -0500 Subject: [PATCH] Implement atexit functionality --- Lib/atexit.py | 9 --------- Lib/logging/__init__.py | 4 ++-- src/main.rs | 39 ++++++++++++++++++++++------------- vm/src/exceptions.rs | 4 ++-- vm/src/lib.rs | 3 +++ vm/src/stdlib/atexit.rs | 45 +++++++++++++++++++++++++++++++++++++++++ vm/src/stdlib/mod.rs | 2 ++ vm/src/sysmodule.rs | 28 +++++++++++++++++++++++++ vm/src/vm.rs | 21 ++++++++++++++++++- 9 files changed, 127 insertions(+), 28 deletions(-) delete mode 100644 Lib/atexit.py create mode 100644 vm/src/stdlib/atexit.rs diff --git a/Lib/atexit.py b/Lib/atexit.py deleted file mode 100644 index c990e82ba..000000000 --- a/Lib/atexit.py +++ /dev/null @@ -1,9 +0,0 @@ -# Dummy implementation of atexit - - -def register(func, *args, **kwargs): - return func - - -def unregister(func): - pass diff --git a/Lib/logging/__init__.py b/Lib/logging/__init__.py index 436087dd7..1a652309a 100644 --- a/Lib/logging/__init__.py +++ b/Lib/logging/__init__.py @@ -2051,8 +2051,8 @@ def shutdown(handlerList=_handlerList): #else, swallow #Let's try and shutdown automatically on application exit... -# import atexit -# atexit.register(shutdown) +import atexit +atexit.register(shutdown) # Null handler diff --git a/src/main.rs b/src/main.rs index 3a3993a76..c059c68b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -53,7 +53,7 @@ fn main() { let interp = Interpreter::new(settings, init); - interp.enter(move |vm| { + let exitcode = interp.enter(move |vm| { let res = run_rustpython(vm, &matches); #[cfg(feature = "flame-it")] @@ -65,22 +65,25 @@ fn main() { } // See if any exception leaked out: - if let Err(err) = res { - if objtype::isinstance(&err, &vm.ctx.exceptions.system_exit) { + let exitcode = match res { + Ok(()) => 0, + Err(err) if objtype::isinstance(&err, &vm.ctx.exceptions.system_exit) => { let args = err.args(); match args.borrow_value().len() { - 0 => return, + 0 => 0, 1 => match_class!(match args.borrow_value()[0].clone() { i @ PyInt => { use num_traits::cast::ToPrimitive; - process::exit(i.borrow_value().to_i32().unwrap_or(0)); + i.borrow_value().to_i32().unwrap_or(0) } arg => { if vm.is_none(&arg) { - return; - } - if let Ok(s) = vm.to_str(&arg) { - eprintln!("{}", s); + 0 + } else { + if let Ok(s) = vm.to_str(&arg) { + eprintln!("{}", s); + } + 1 } } }), @@ -88,14 +91,22 @@ fn main() { if let Ok(r) = vm.to_repr(args.as_object()) { eprintln!("{}", r); } + 1 } } - } else { - print_exception(&vm, err); } - process::exit(1); - } - }) + Err(err) => { + print_exception(&vm, err); + 1 + } + }; + + let _ = vm.run_atexit_funcs(); + + exitcode + }); + + process::exit(exitcode); } fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> { diff --git a/vm/src/exceptions.rs b/vm/src/exceptions.rs index 1cab3e718..d2cb2f766 100644 --- a/vm/src/exceptions.rs +++ b/vm/src/exceptions.rs @@ -162,10 +162,10 @@ pub fn print_exception(vm: &VirtualMachine, exc: PyBaseExceptionRef) { } }; if let Ok(excepthook) = vm.get_attribute(vm.sys_module.clone(), "excepthook") { - let (exc_type, exc_val, exc_tb) = split(exc, vm); + let (exc_type, exc_val, exc_tb) = split(exc.clone(), vm); if let Err(eh_exc) = vm.invoke(&excepthook, vec![exc_type, exc_val, exc_tb]) { write_fallback(&eh_exc, "Error in sys.excepthook:"); - write_fallback(&eh_exc, "Original exception was:"); + write_fallback(&exc, "Original exception was:"); } } else { write_fallback(&exc, "missing sys.excepthook"); diff --git a/vm/src/lib.rs b/vm/src/lib.rs index ad4e71edb..3354d7764 100644 --- a/vm/src/lib.rs +++ b/vm/src/lib.rs @@ -7,6 +7,9 @@ // for methods like vm.to_str(), not the typical use of 'to' as a method prefix #![allow(clippy::wrong_self_convention, clippy::implicit_hasher)] +// to allow `mod foo {}` in foo.rs; clippy thinks this is a mistake/misunderstanding of +// how `mod` works, but we want this sometimes for pymodule declarations +#![allow(clippy::module_inception)] #![doc(html_logo_url = "https://raw.githubusercontent.com/RustPython/RustPython/master/logo.png")] #![doc(html_root_url = "https://docs.rs/rustpython-vm/")] #![cfg_attr( diff --git a/vm/src/stdlib/atexit.rs b/vm/src/stdlib/atexit.rs new file mode 100644 index 000000000..4a18d49ab --- /dev/null +++ b/vm/src/stdlib/atexit.rs @@ -0,0 +1,45 @@ +pub(crate) use atexit::make_module; + +#[pymodule] +mod atexit { + use crate::function::PyFuncArgs; + use crate::pyobject::{PyObjectRef, PyResult}; + use crate::VirtualMachine; + + #[pyfunction] + fn register(func: PyObjectRef, args: PyFuncArgs, vm: &VirtualMachine) -> PyObjectRef { + vm.state.atexit_funcs.lock().push((func.clone(), args)); + func + } + + #[pyfunction] + fn _clear(vm: &VirtualMachine) { + vm.state.atexit_funcs.lock().clear(); + } + + #[pyfunction] + fn unregister(func: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let mut funcs = vm.state.atexit_funcs.lock(); + + let mut i = 0; + while i < funcs.len() { + if vm.bool_eq(funcs[i].0.clone(), func.clone())? { + funcs.remove(i); + } else { + i += 1; + } + } + + Ok(()) + } + + #[pyfunction] + fn _run_exitfuncs(vm: &VirtualMachine) -> PyResult<()> { + vm.run_atexit_funcs() + } + + #[pyfunction] + fn _ncallbacks(vm: &VirtualMachine) -> usize { + vm.state.atexit_funcs.lock().len() + } +} diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs index 3df65bc0f..05ef201ed 100644 --- a/vm/src/stdlib/mod.rs +++ b/vm/src/stdlib/mod.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; pub mod array; #[cfg(feature = "rustpython-parser")] pub(crate) mod ast; +mod atexit; mod binascii; mod collections; mod csv; @@ -73,6 +74,7 @@ pub fn get_module_inits() -> HashMap { #[allow(unused_mut)] let mut modules = hashmap! { "array".to_owned() => Box::new(array::make_module) as StdlibInitFunc, + "atexit".to_owned() => Box::new(atexit::make_module), "binascii".to_owned() => Box::new(binascii::make_module), "_collections".to_owned() => Box::new(collections::make_module), "_csv".to_owned() => Box::new(csv::make_module), diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs index dc7d55c62..f8dff4238 100644 --- a/vm/src/sysmodule.rs +++ b/vm/src/sysmodule.rs @@ -332,6 +332,34 @@ pub fn get_stderr(vm: &VirtualMachine) -> PyResult { .map_err(|_| vm.new_runtime_error("lost sys.stderr".to_owned())) } +/// Similar to PySys_WriteStderr in CPython. +/// +/// # Usage +/// +/// ```rust,ignore +/// writeln!(sysmodule::PyStderr(vm), "foo bar baz :)"); +/// ``` +/// +/// Unlike writing to a `std::io::Write` with the `write[ln]!()` macro, there's no error condition here; +/// this is intended to be a replacement for the `eprint[ln]!()` macro, so `write!()`-ing to PyStderr just +/// returns `()`. +pub struct PyStderr<'vm>(pub &'vm VirtualMachine); + +impl PyStderr<'_> { + pub fn write_fmt(&self, args: std::fmt::Arguments<'_>) { + use py_io::Write; + + let vm = self.0; + if let Ok(stderr) = get_stderr(vm) { + let mut stderr = py_io::PyWriter(stderr, vm); + if let Ok(()) = stderr.write_fmt(args) { + return; + } + } + eprint!("{}", args) + } +} + fn sys_excepthook( exc_type: PyObjectRef, exc_val: PyObjectRef, diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 90d191f64..a3357540d 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -19,7 +19,7 @@ use rustpython_compiler::{ use crate::builtins::{self, to_ascii}; use crate::bytecode; -use crate::common::{hash::HashSecret, rc::PyRc}; +use crate::common::{cell::PyMutex, hash::HashSecret, rc::PyRc}; use crate::exceptions::{self, PyBaseException, PyBaseExceptionRef}; use crate::frame::{ExecutionResult, Frame, FrameRef}; use crate::frozen; @@ -122,6 +122,7 @@ pub struct PyGlobalState { pub stacksize: AtomicCell, pub thread_count: AtomicCell, pub hash_secret: HashSecret, + pub atexit_funcs: PyMutex>, } pub const NSIG: usize = 64; @@ -258,6 +259,7 @@ impl VirtualMachine { stacksize: AtomicCell::new(0), thread_count: AtomicCell::new(0), hash_secret, + atexit_funcs: PyMutex::default(), }), initialized: false, }; @@ -344,6 +346,23 @@ impl VirtualMachine { } } + pub fn run_atexit_funcs(&self) -> PyResult<()> { + let mut last_exc = None; + for (func, args) in self.state.atexit_funcs.lock().drain(..).rev() { + if let Err(e) = self.invoke(&func, args) { + last_exc = Some(e.clone()); + if !objtype::isinstance(&e, &self.ctx.exceptions.system_exit) { + writeln!(sysmodule::PyStderr(self), "Error in atexit._run_exitfuncs:"); + exceptions::print_exception(self, e); + } + } + } + match last_exc { + None => Ok(()), + Some(e) => Err(e), + } + } + pub fn run_code_obj(&self, code: PyCodeRef, scope: Scope) -> PyResult { let frame = Frame::new(code, scope, self).into_ref(self); self.run_frame_full(frame)