Fixes #250: exit() now raises SystemExit

Signed-off-by: Janos <86970079+janosdebugs@users.noreply.github.com>

include/pocketpy/vm.h

@@ -242,7 +242,27 @@ public:
     template<typename ...Args>
     PyVar _exec(Args&&... args){
         callstack.emplace(s_data._sp, std::forward<Args>(args)...);
-        return __run_top_frame();
+        try {
+            return __run_top_frame();
+        } catch (Exception& e) {
+            // TODO Checking the exception type by name seems ugly and probably won't work on derived classes.
+            //      There must be a better way to do this?
+            // TODO Clean up the nesting here.
+            if (e.type == "SystemExit") {
+                i64 exit_code;
+                // Make sure the variable is preserved after the SystemExit is gc'd.
+                // TODO: is this the right way to do it?
+                if (try_cast_int(e._self->attr("code"), &exit_code)) {
+                    // TODO Copied from py_var below:
+                    i64 val = static_cast<i64>(std::forward<i64>(exit_code));
+                    if(val >= Number::kMinSmallInt && val <= Number::kMaxSmallInt){
+                        val = (val << 2) | 0b10;
+                        return reinterpret_cast<PyVar>(val);
+                    }
+                }
+            }
+            throw;
+        }
     }
 #endif
 
@@ -357,6 +377,7 @@ public:
     void BinaryOptError(const char* op, PyVar _0, PyVar _1);
     void AttributeError(PyVar obj, StrName name);
     void AttributeError(const Str& msg){ __builtin_error("AttributeError", msg); }
+    void SystemExit(const int code){ __builtin_error("SystemExit", code); }
 #endif
 
 #if PK_REGION("Type Checking Methods")
@@ -429,6 +450,7 @@ public:
     void __builtin_error(StrName type);
     void __builtin_error(StrName type, PyVar arg);
     void __builtin_error(StrName type, const Str& msg);
+    void __builtin_error(StrName type, const int code);
     void __push_varargs(){}
     void __push_varargs(PyVar _0){ PUSH(_0); }
     void __push_varargs(PyVar _0, PyVar _1){ PUSH(_0); PUSH(_1); }

python/builtins.py

@@ -214,3 +214,7 @@ class KeyError(Exception):
         if self.key is ...:
             return 'KeyError()'
         return f'KeyError({self.key!r})'
+
+class SystemExit(Exception):
+    def __init__(self, code=None):
+        self.code = code

src/pocketpy.cpp

@@ -224,7 +224,7 @@ void __init_builtins(VM* _vm) {
     });
 
     _vm->bind(_vm->builtins, "exit(code=0)", [](VM* vm, ArgsView args) {
-        std::exit(CAST(int, args[0]));
+        vm->SystemExit(CAST(int, args[0]));
         return vm->None;
     });
 

src/vm.cpp

@@ -1373,6 +1373,7 @@ PyVar VM::bind_property(PyVar obj, const char* name, NativeFuncC fget, NativeFun
 void VM::__builtin_error(StrName type){ _error(call(builtins->attr(type))); }
 void VM::__builtin_error(StrName type, PyVar arg){ _error(call(builtins->attr(type), arg)); }
 void VM::__builtin_error(StrName type, const Str& msg){ __builtin_error(type, VAR(msg)); }
+void VM::__builtin_error(StrName type, const int code){ __builtin_error(type, VAR(code)); }
 
 void VM::BinaryOptError(const char* op, PyVar _0, PyVar _1) {
     StrName name_0 = _type_name(vm, _tp(_0));