exception-to-exitcode to vm method

Author: youknowone

src/lib.rs

@@ -47,11 +47,9 @@ 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,
+    Interpreter, PyResult, Settings, VirtualMachine,
 };
 use std::{env, process, str::FromStr};
 
@@ -68,7 +66,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)]
@@ -89,61 +88,31 @@ where
     });
 
     let exitcode = interp.enter(move |vm| {
-        let res = run_rustpython(vm, &matches);
+        let res = 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);
-            }
-        }
-
         // 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();
-                let exitcode = match args.as_slice() {
-                    [] => Ok(0),
-                    [arg] => match_class!(match arg {
-                        ref i @ PyInt => {
-                            use num_traits::cast::ToPrimitive;
-                            Ok(i.as_bigint().to_i32().unwrap_or(0))
-                        }
-                        arg => {
-                            if vm.is_none(arg) {
-                                Ok(0)
-                            } else {
-                                Err(arg.str(vm).ok())
-                            }
-                        }
-                    }),
-                    _ => Err(args.as_object().repr(vm).ok()),
-                };
-                exitcode.unwrap_or_else(|msg| {
-                    if let Some(msg) = msg {
-                        let stderr = sys::PyStderr(vm);
-                        writeln!(stderr, "{}", msg);
-                    }
-                    1
-                })
-            }
-            Err(exc) => {
-                vm.print_exception(exc);
-                1
-            }
-        };
+        let exit_code = res
+            .map(|_| 0)
+            .map_err(|exc| vm.handle_exit_exception(exc))
+            .unwrap_or_else(|code| code);
 
         let _ = atexit::_run_exitfuncs(vm);
 
         flush_std(vm);
 
-        exitcode
+        exit_code
     });
 
+    #[cfg(feature = "flame-it")]
+    {
+        main_guard.end();
+        if let Err(e) = write_profile(&matches) {
+            error!("Error writing profile information: {}", e);
+        }
+    }
+
     process::exit(exitcode)
 }
 

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))
     }