Correctly handle a SystemExit unwound to the top level of execution

coolreader18 · coolreader18 · commit e89d00f16718 · 2019-10-04T22:05:01.000-05:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,6 +27,7 @@ rustpython-compiler = {path = "compiler", version = "0.1.1"}
 rustpython-parser = {path = "parser", version = "0.1.1"}
 rustpython-vm = {path = "vm", version = "0.1.1"}
 dirs = "2.0"
+num-traits = "0.2.8"
 
 flame = { version = "0.2", optional = true }
 flamescope = { version = "0.1", optional = true }
diff --git a/src/main.rs b/src/main.rs
@@ -8,7 +8,9 @@ use clap::{App, AppSettings, Arg, ArgMatches};
 use rustpython_compiler::{compile, error::CompileError, error::CompileErrorType};
 use rustpython_parser::error::ParseErrorType;
 use rustpython_vm::{
-    import, print_exception,
+    import, match_class,
+    obj::{objint::PyInt, objtuple::PyTuple, objtype},
+    print_exception,
     pyobject::{ItemProtocol, PyObjectRef, PyResult},
     scope::Scope,
     util, PySettings, VirtualMachine,
@@ -30,17 +32,47 @@ fn main() {
     let vm = VirtualMachine::new(settings);
 
     let res = run_rustpython(&vm, &matches);
-    // See if any exception leaked out:
-    handle_exception(&vm, res);
 
     #[cfg(feature = "flame-it")]
     {
         main_guard.end();
         if let Err(e) = write_profile(&matches) {
             error!("Error writing profile information: {}", e);
-            process::exit(1);
         }
     }
+
+    // See if any exception leaked out:
+    if let Err(err) = res {
+        if objtype::isinstance(&err, &vm.ctx.exceptions.system_exit) {
+            let args = vm.get_attribute(err.clone(), "args").unwrap();
+            let args = args.downcast::<PyTuple>().expect("'args' must be a tuple");
+            match args.elements.len() {
+                0 => return,
+                1 => match_class!(match args.elements[0].clone() {
+                    i @ PyInt => {
+                        use num_traits::cast::ToPrimitive;
+                        process::exit(i.as_bigint().to_i32().unwrap());
+                    }
+                    arg => {
+                        if vm.is_none(&arg) {
+                            return;
+                        }
+                        if let Ok(s) = vm.to_str(&arg) {
+                            println!("{}", s);
+                        }
+                    }
+                }),
+                _ => {
+                    if let Ok(r) = vm.to_repr(args.as_object()) {
+                        println!("{}", r);
+                    }
+                }
+            }
+        } else {
+            print_exception(&vm, &err);
+        }
+        process::exit(1);
+    }
 }
 
 fn parse_arguments<'a>(app: App<'a, '_>) -> ArgMatches<'a> {
@@ -349,13 +381,6 @@ fn _run_string(vm: &VirtualMachine, scope: Scope, source: &str, source_path: Str
     vm.run_code_obj(code_obj, scope)
 }
 
-fn handle_exception<T>(vm: &VirtualMachine, result: PyResult<T>) {
-    if let Err(err) = result {
-        print_exception(vm, &err);
-        process::exit(1);
-    }
-}
-
 fn run_command(vm: &VirtualMachine, scope: Scope, source: String) -> PyResult<()> {
     debug!("Running command {}", source);
     _run_string(vm, scope, &source, "<stdin>".to_string())?;
@@ -560,6 +585,10 @@ fn run_shell(vm: &VirtualMachine, scope: Scope) -> PyResult<()> {
         };
 
         if let Err(exc) = result {
+            if objtype::isinstance(&exc, &vm.ctx.exceptions.system_exit) {
+                repl.save_history(&repl_history_path).unwrap();
+                return Err(exc);
+            }
             print_exception(vm, &exc);
         }
     }
diff --git a/vm/src/sysmodule.rs b/vm/src/sysmodule.rs
@@ -171,10 +171,9 @@ fn sys_git_info(vm: &VirtualMachine) -> PyObjectRef {
     ])
 }
 
-// TODO: raise a SystemExit here
-fn sys_exit(code: OptionalArg<i32>, _vm: &VirtualMachine) -> PyResult<()> {
-    let code = code.unwrap_or(0);
-    std::process::exit(code)
+fn sys_exit(code: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult {
+    let code = code.unwrap_or_else(|| vm.new_int(0));
+    Err(vm.new_exception_obj(vm.ctx.exceptions.system_exit.clone(), vec![code])?)
 }
 
 pub fn make_module(vm: &VirtualMachine, module: PyObjectRef, builtins: PyObjectRef) {
diff --git a/vm/src/vm.rs b/vm/src/vm.rs
@@ -303,7 +303,7 @@ impl VirtualMachine {
     }
 
     #[cfg_attr(feature = "flame-it", flame("VirtualMachine"))]
-    fn new_exception_obj(&self, exc_type: PyClassRef, args: Vec<PyObjectRef>) -> PyResult {
+    pub fn new_exception_obj(&self, exc_type: PyClassRef, args: Vec<PyObjectRef>) -> PyResult {
         // TODO: add repr of args into logging?
         vm_trace!("New exception created: {}", exc_type.name);
         self.invoke(&exc_type.into_object(), args)

Original file line number	Diff line number	Diff line change
`@@ -171,10 +171,9 @@ fn sys_git_info(vm: &VirtualMachine) -> PyObjectRef {`
`171`	`171`	`])`
`172`	`172`	`}`
`173`	`173`
`174`		`-// TODO: raise a SystemExit here`
`175`		`-fn sys_exit(code: OptionalArg<i32>, _vm: &VirtualMachine) -> PyResult<()> {`
`176`		`- let code = code.unwrap_or(0);`
`177`		`- std::process::exit(code)`
	`174`	`+fn sys_exit(code: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult {`
	`175`	`+ let code = code.unwrap_or_else(\|\| vm.new_int(0));`
	`176`	`+ Err(vm.new_exception_obj(vm.ctx.exceptions.system_exit.clone(), vec![code])?)`
`178`	`177`	`}`
`179`	`178`
`180`	`179`	`pub fn make_module(vm: &VirtualMachine, module: PyObjectRef, builtins: PyObjectRef) {`
Original file line number	Diff line number	Diff line change
`@@ -303,7 +303,7 @@ impl VirtualMachine {`
`303`	`303`	`}`
`304`	`304`
`305`	`305`	`#[cfg_attr(feature = "flame-it", flame("VirtualMachine"))]`
`306`		`- fn new_exception_obj(&self, exc_type: PyClassRef, args: Vec<PyObjectRef>) -> PyResult {`
	`306`	`+ pub fn new_exception_obj(&self, exc_type: PyClassRef, args: Vec<PyObjectRef>) -> PyResult {`
`307`	`307`	`// TODO: add repr of args into logging?`
`308`	`308`	`vm_trace!("New exception created: {}", exc_type.name);`
`309`	`309`	`self.invoke(&exc_type.into_object(), args)`