Implement atexit functionality

Author: coolreader18

Lib/atexit.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
 

Lib/logging/__init__.py

@@ -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,37 +65,48 @@ 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
                             }
                         }
                     }),
                     _ => {
                         if let Ok(r) = vm.to_repr(args.as_object()) {
                             eprintln!("{}", r);
                         }
+                        1
                     }
                 }
-            } else {
+            }
+            Err(err) => {
                 print_exception(&vm, err);
+                1
             }
-            process::exit(1);
-        }
-    })
+        };
+
+        let _ = vm.run_atexit_funcs();
+
+        exitcode
+    });
+
+    process::exit(exitcode);
 }
 
 fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> {

src/main.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");

vm/src/exceptions.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(

vm/src/lib.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()
+    }
+}

vm/src/stdlib/atexit.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<String, StdlibInitFunc> {
     #[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),

vm/src/stdlib/mod.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,

vm/src/sysmodule.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<usize>,
     pub thread_count: AtomicCell<usize>,
     pub hash_secret: HashSecret,
+    pub atexit_funcs: PyMutex<Vec<(PyObjectRef, PyFuncArgs)>>,
 }
 
 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)