Merge pull request RustPython#3672 from youknowone/exit-handling

Fix stderr usage in exit handling

Author: coolreader18

src/lib.rs

@@ -46,13 +46,7 @@ extern crate log;
 mod shell;
 
 use clap::{App, AppSettings, Arg, ArgMatches};
-use rustpython_vm::{
-    builtins::PyInt,
-    match_class,
-    scope::Scope,
-    stdlib::{atexit, sys},
-    AsObject, Interpreter, PyResult, Settings, VirtualMachine,
-};
+use rustpython_vm::{scope::Scope, Interpreter, PyResult, Settings, VirtualMachine};
 use std::{env, process, str::FromStr};
 
 pub use rustpython_vm as vm;
@@ -68,7 +62,8 @@ where
     env_logger::init();
     let app = App::new("RustPython");
     let matches = parse_arguments(app);
-    let settings = create_settings(&matches);
+    let matches = &matches;
+    let settings = create_settings(matches);
 
     // don't translate newlines (\r\n <=> \n)
     #[cfg(windows)]
@@ -88,75 +83,19 @@ where
         init(vm);
     });
 
-    let exitcode = interp.enter(move |vm| {
-        let res = run_rustpython(vm, &matches);
+    let exitcode = interp.run(move |vm| run_rustpython(vm, matches));
 
-        flush_std(vm);
-
-        #[cfg(feature = "flame-it")]
-        {
-            main_guard.end();
-            if let Err(e) = write_profile(&matches) {
-                error!("Error writing profile information: {}", e);
-            }
+    #[cfg(feature = "flame-it")]
+    {
+        main_guard.end();
+        if let Err(e) = write_profile(&matches) {
+            error!("Error writing profile information: {}", e);
         }
-
-        // See if any exception leaked out:
-        let exitcode = match res {
-            Ok(()) => 0,
-            Err(err) if err.fast_isinstance(&vm.ctx.exceptions.system_exit) => {
-                let args = err.args();
-                match args.as_slice() {
-                    [] => 0,
-                    [arg] => match_class!(match arg {
-                        ref i @ PyInt => {
-                            use num_traits::cast::ToPrimitive;
-                            i.as_bigint().to_i32().unwrap_or(0)
-                        }
-                        arg => {
-                            if vm.is_none(arg) {
-                                0
-                            } else {
-                                if let Ok(s) = arg.str(vm) {
-                                    eprintln!("{}", s);
-                                }
-                                1
-                            }
-                        }
-                    }),
-                    _ => {
-                        if let Ok(r) = args.as_object().repr(vm) {
-                            eprintln!("{}", r);
-                        }
-                        1
-                    }
-                }
-            }
-            Err(exc) => {
-                vm.print_exception(exc);
-                1
-            }
-        };
-
-        let _ = atexit::_run_exitfuncs(vm);
-
-        flush_std(vm);
-
-        exitcode
-    });
+    }
 
     process::exit(exitcode)
 }
 
-fn flush_std(vm: &VirtualMachine) {
-    if let Ok(stdout) = sys::get_stdout(vm) {
-        let _ = vm.call_method(&stdout, "flush", ());
-    }
-    if let Ok(stderr) = sys::get_stderr(vm) {
-        let _ = vm.call_method(&stderr, "flush", ());
-    }
-}
-
 fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> {
     let app = app
         .setting(AppSettings::TrailingVarArg)

stdlib/src/faulthandler.rs

@@ -2,10 +2,14 @@ pub(crate) use decl::make_module;
 
 #[pymodule(name = "faulthandler")]
 mod decl {
-    use crate::vm::{frame::FrameRef, function::OptionalArg, VirtualMachine};
+    use crate::vm::{
+        frame::FrameRef, function::OptionalArg, stdlib::sys::PyStderr, VirtualMachine,
+    };
 
-    fn dump_frame(frame: &FrameRef) {
-        eprintln!(
+    fn dump_frame(frame: &FrameRef, vm: &VirtualMachine) {
+        let stderr = PyStderr(vm);
+        writeln!(
+            stderr,
             "  File \"{}\", line {} in {}",
             frame.code.source_path,
             frame.current_location().row(),
@@ -19,10 +23,11 @@ mod decl {
         _all_threads: OptionalArg<bool>,
         vm: &VirtualMachine,
     ) {
-        eprintln!("Stack (most recent call first):");
+        let stderr = PyStderr(vm);
+        writeln!(stderr, "Stack (most recent call first):");
 
         for frame in vm.frames.borrow().iter() {
-            dump_frame(frame);
+            dump_frame(frame, vm);
         }
     }
 

vm/src/stdlib/sys.rs

@@ -501,7 +501,8 @@ mod sys {
     #[pyfunction(name = "__unraisablehook__")]
     fn unraisablehook(unraisable: UnraisableHookArgs, vm: &VirtualMachine) {
         if let Err(e) = _unraisablehook(unraisable, vm) {
-            println!("{}", e.as_object().repr(vm).unwrap().as_str());
+            let stderr = super::PyStderr(vm);
+            writeln!(stderr, "{}", e.as_object().repr(vm).unwrap().as_str());
         }
     }
 

vm/src/stdlib/warnings.rs

@@ -5,6 +5,7 @@ mod _warnings {
     use crate::{
         builtins::{PyStrRef, PyTypeRef},
         function::OptionalArg,
+        stdlib::sys::PyStderr,
         AsObject, PyResult, VirtualMachine,
     };
 
@@ -33,7 +34,14 @@ mod _warnings {
         } else {
             vm.ctx.exceptions.user_warning.clone()
         };
-        eprintln!("level:{}: {}: {}", level, category.name(), args.message);
+        let stderr = PyStderr(vm);
+        writeln!(
+            stderr,
+            "level:{}: {}: {}",
+            level,
+            category.name(),
+            args.message
+        );
         Ok(())
     }
 }

vm/src/vm/interpreter.rs

@@ -1,4 +1,8 @@
 use super::{setting::Settings, thread, VirtualMachine};
+use crate::{
+    stdlib::{atexit, sys},
+    PyResult,
+};
 
 /// The general interface for the VM
 ///
@@ -53,16 +57,36 @@ impl Interpreter {
         thread::enter_vm(&self.vm, || f(&self.vm))
     }
 
-    // TODO: interpreter shutdown
-    // pub fn run<F>(self, f: F)
-    // where
-    //     F: FnOnce(&VirtualMachine),
-    // {
-    //     self.enter(f);
-    //     self.shutdown();
-    // }
+    pub fn run<F, R>(self, f: F) -> i32
+    where
+        F: FnOnce(&VirtualMachine) -> PyResult<R>,
+    {
+        self.enter(|vm| {
+            let res = f(vm);
+            flush_std(vm);
+
+            // See if any exception leaked out:
+            let exit_code = res
+                .map(|_| 0)
+                .map_err(|exc| vm.handle_exit_exception(exc))
+                .unwrap_or_else(|code| code);
 
-    // pub fn shutdown(self) {}
+            let _ = atexit::_run_exitfuncs(vm);
+
+            flush_std(vm);
+
+            exit_code
+        })
+    }
+}
+
+fn flush_std(vm: &VirtualMachine) {
+    if let Ok(stdout) = sys::get_stdout(vm) {
+        let _ = vm.call_method(&stdout, "flush", ());
+    }
+    if let Ok(stderr) = sys::get_stderr(vm) {
+        let _ = vm.call_method(&stderr, "flush", ());
+    }
 }
 
 #[cfg(test)]

vm/src/vm/mod.rs

@@ -19,7 +19,7 @@ use crate::{
         code::{self, PyCode},
         pystr::IntoPyStrRef,
         tuple::{PyTuple, PyTupleTyped},
-        PyBaseExceptionRef, PyDictRef, PyList, PyModule, PyStrRef, PyTypeRef,
+        PyBaseExceptionRef, PyDictRef, PyInt, PyList, PyModule, PyStrRef, PyTypeRef,
     },
     bytecode,
     codecs::CodecsRegistry,
@@ -278,15 +278,13 @@ impl VirtualMachine {
 
     #[cold]
     pub fn run_unraisable(&self, e: PyBaseExceptionRef, msg: Option<String>, object: PyObjectRef) {
-        use crate::stdlib::sys::UnraisableHookArgs;
-
         let sys_module = self.import("sys", None, 0).unwrap();
         let unraisablehook = sys_module.get_attr("unraisablehook", self).unwrap();
 
         let exc_type = e.class().clone();
         let exc_traceback = e.traceback().to_pyobject(self); // TODO: actual traceback
         let exc_value = e.into();
-        let args = UnraisableHookArgs {
+        let args = stdlib::sys::UnraisableHookArgs {
             exc_type,
             exc_value,
             exc_traceback,
@@ -660,6 +658,37 @@ impl VirtualMachine {
         }
     }
 
+    pub fn handle_exit_exception(&self, exc: PyBaseExceptionRef) -> i32 {
+        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;
+                        return i.as_bigint().to_i32().unwrap_or(0);
+                    }
+                    arg => {
+                        if self.is_none(arg) {
+                            return 0;
+                        } else {
+                            arg.str(self).ok()
+                        }
+                    }
+                }),
+                _ => args.as_object().repr(self).ok(),
+            };
+            if let Some(msg) = msg {
+                let stderr = stdlib::sys::PyStderr(self);
+                writeln!(stderr, "{}", msg);
+            }
+            1
+        } else {
+            self.print_exception(exc);
+            1
+        }
+    }
+
     pub fn map_codeobj(&self, code: bytecode::CodeObject) -> code::CodeObject {
         code.map_bag(&code::PyObjBag(self))
     }