From e89d00f167187f5514a75910285f6901e78beadf Mon Sep 17 00:00:00 2001
From: Noah <33094578+coolreader18@users.noreply.github.com>
Date: Thu, 3 Oct 2019 19:16:00 +0000
Subject: [PATCH] Correctly handle a SystemExit unwound to the top level of
 execution

---
 Cargo.lock          |  1 +
 Cargo.toml          |  1 +
 src/main.rs         | 51 +++++++++++++++++++++++++++++++++++----------
 vm/src/sysmodule.rs |  7 +++----
 vm/src/vm.rs        |  2 +-
 5 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 228def36d..5e845c408 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1092,6 +1092,7 @@ dependencies = [
  "flame 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "flamescope 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)",
+ "num-traits 0.2.8 (registry+https://github.com/rust-lang/crates.io-index)",
  "rustpython-compiler 0.1.1",
  "rustpython-parser 0.1.1",
  "rustpython-vm 0.1.1",
diff --git a/Cargo.toml b/Cargo.toml
index 2a9d826ab..8192b42a2 100644
--- 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
index 104639b3a..65adc69e4 100644
--- 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
index 8abc12558..e695526cf 100644
--- 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
index 4858daeab..27e8cf05e 100644
--- 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)