From 5d4294b457589575671b623584abdca1670e2d7f Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 14 Apr 2024 13:29:20 +0800
Subject: [PATCH] add `pkpy_set_main_argv`

---
 docs/C-API/introduction.md    |  4 ++++
 include/pocketpy/frame.h      |  2 +-
 include/pocketpy/obj.h        |  2 +-
 include/pocketpy/pocketpy_c.h |  1 +
 include/pocketpy/vm.h         | 12 ++++++++++--
 src/pocketpy_c.cpp            |  5 +++++
 src/vm.cpp                    | 15 +++++----------
 src2/main.cpp                 |  2 ++
 src2/pocketpy_c.c             |  4 ++++
 tests/87_sys.py               |  4 ++++
 10 files changed, 37 insertions(+), 14 deletions(-)
 create mode 100644 tests/87_sys.py

diff --git a/docs/C-API/introduction.md b/docs/C-API/introduction.md
index 1351be9f..397e8745 100644
--- a/docs/C-API/introduction.md
+++ b/docs/C-API/introduction.md
@@ -36,3 +36,7 @@ C-APIs are always stable and backward compatible.
 + `bool pkpy_exec_2(pkpy_vm*, const char* source, const char* filename, int mode, const char* module)`
 
     Wraps `vm->exec_2`. Execute a string of source code with more options.
+
++ `void pkpy_set_main_argv(pkpy_vm*, int argc, char** argv)`
+
+    Wraps `vm->set_main_argv`. Set the `sys.argv` before executing scripts.
diff --git a/include/pocketpy/frame.h b/include/pocketpy/frame.h
index e1466273..0b322171 100644
--- a/include/pocketpy/frame.h
+++ b/include/pocketpy/frame.h
@@ -84,7 +84,7 @@ struct Frame {
     PyObject* _callable;    // a function object or nullptr (global scope)
     FastLocals _locals;
 
-    NameDict& f_globals() noexcept { return _module->attr(); }
+    NameDict& f_globals() { return _module->attr(); }
     PyObject* f_closure_try_get(StrName name);
 
     // function scope
diff --git a/include/pocketpy/obj.h b/include/pocketpy/obj.h
index a1d5f8e7..8391fa2b 100644
--- a/include/pocketpy/obj.h
+++ b/include/pocketpy/obj.h
@@ -180,7 +180,7 @@ struct Py_ final: PyObject {
 struct MappingProxy{
     PyObject* obj;
     MappingProxy(PyObject* obj) : obj(obj) {}
-    NameDict& attr() noexcept { return obj->attr(); }
+    NameDict& attr() { return obj->attr(); }
 };
 
 #define PK_OBJ_GET(T, obj) (((Py_<T>*)(obj))->_value)
diff --git a/include/pocketpy/pocketpy_c.h b/include/pocketpy/pocketpy_c.h
index ba3d8dc2..f5c8a000 100644
--- a/include/pocketpy/pocketpy_c.h
+++ b/include/pocketpy/pocketpy_c.h
@@ -23,6 +23,7 @@ PK_EXPORT pkpy_vm* pkpy_new_vm(bool enable_os);
 PK_EXPORT void pkpy_delete_vm(pkpy_vm*);
 PK_EXPORT bool pkpy_exec(pkpy_vm*, const char* source);
 PK_EXPORT bool pkpy_exec_2(pkpy_vm*, const char* source, const char* filename, int mode, const char* module);
+PK_EXPORT void pkpy_set_main_argv(pkpy_vm*, int argc, char** argv);
 
 /* Stack Manipulation */
 PK_EXPORT bool pkpy_dup(pkpy_vm*, int i);
diff --git a/include/pocketpy/vm.h b/include/pocketpy/vm.h
index de5cf27f..9ba5ee55 100644
--- a/include/pocketpy/vm.h
+++ b/include/pocketpy/vm.h
@@ -153,8 +153,16 @@ public:
 
     VM(bool enable_os=true);
 
-    Frame* top_frame();
-    void _pop_frame();
+    void set_main_argv(int argc, char** argv);
+
+    Frame* top_frame(){
+        return &callstack.top();
+    }
+
+    void _pop_frame(){
+        s_data.reset(callstack.top()._sp_base);
+        callstack.pop();
+    }
 
     PyObject* py_str(PyObject* obj);
     PyObject* py_repr(PyObject* obj);
diff --git a/src/pocketpy_c.cpp b/src/pocketpy_c.cpp
index 1311cb9a..47ce89ae 100644
--- a/src/pocketpy_c.cpp
+++ b/src/pocketpy_c.cpp
@@ -91,6 +91,11 @@ bool pkpy_exec_2(pkpy_vm* vm_handle, const char* source, const char* filename, i
     return res != nullptr;
 }
 
+void pkpy_set_main_argv(pkpy_vm* vm_handle, int argc, char** argv){
+    VM* vm = (VM*) vm_handle;
+    vm->set_main_argv(argc, argv);
+}
+
 bool pkpy_dup(pkpy_vm* vm_handle, int n){
     VM* vm = (VM*) vm_handle;
     PK_ASSERT_NO_ERROR()
diff --git a/src/vm.cpp b/src/vm.cpp
index 383de456..b22895a8 100644
--- a/src/vm.cpp
+++ b/src/vm.cpp
@@ -119,16 +119,11 @@ namespace pkpy{
         PK_UNREACHABLE();
     }
 
-    Frame* VM::top_frame(){
-#if PK_DEBUG_EXTRA_CHECK
-        if(callstack.empty()) PK_FATAL_ERROR();
-#endif
-        return &callstack.top();
-    }
-
-    void VM::_pop_frame(){
-        s_data.reset(callstack.top()._sp_base);
-        callstack.pop();
+    void VM::set_main_argv(int argc, char** argv){
+        PyObject* mod = vm->_modules["sys"];
+        List argv_(argc);
+        for(int i=0; i<argc; i++) argv_[i] = VAR(std::string_view(argv[i]));
+        mod->attr().set("argv", VAR(std::move(argv_)));
     }
 
     PyObject* VM::find_name_in_mro(Type cls, StrName name){
diff --git a/src2/main.cpp b/src2/main.cpp
index 7830c97c..f0a93cfe 100644
--- a/src2/main.cpp
+++ b/src2/main.cpp
@@ -99,6 +99,8 @@ int main(int argc, char** argv){
         std::string src((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
         file.close();
 
+        pkpy_set_main_argv(vm, argc, argv);
+
         bool ok = pkpy_exec_2(vm, src.c_str(), filepath.filename().string().c_str(), 0, NULL);
         if(!ok) pkpy_clear_error(vm, NULL);
         pkpy_delete_vm(vm);
diff --git a/src2/pocketpy_c.c b/src2/pocketpy_c.c
index 228fc8d4..74727738 100644
--- a/src2/pocketpy_c.c
+++ b/src2/pocketpy_c.c
@@ -24,6 +24,10 @@ bool pkpy_exec_2(pkpy_vm* vm, const char* source, const char* filename, int mode
     return returnValue;
 }
 
+void pkpy_set_main_argv(pkpy_vm* vm, int argc, char** argv){
+
+}
+
 bool pkpy_dup(pkpy_vm* vm, int i) {
     bool returnValue;
     return returnValue;
diff --git a/tests/87_sys.py b/tests/87_sys.py
new file mode 100644
index 00000000..c09d2fd0
--- /dev/null
+++ b/tests/87_sys.py
@@ -0,0 +1,4 @@
+import sys
+
+assert len(sys.argv) == 2
+assert sys.argv[1] == 'tests/87_sys.py'