From 2336d22666c03532889c122e30dd8835b438ce13 Mon Sep 17 00:00:00 2001
From: liuyang18 <blueloveTH@foxmail.com>
Date: Mon, 24 Mar 2025 11:11:03 +0800
Subject: [PATCH] refactor `globals` and `locals`

---
 include/pocketpy/pocketpy.h |  4 ++++
 src/interpreter/frame.c     | 29 +++++++++++++++++++++++++++++
 src/public/modules.c        | 25 ++-----------------------
 3 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h
index 86a11547..8215a0cb 100644
--- a/include/pocketpy/pocketpy.h
+++ b/include/pocketpy/pocketpy.h
@@ -110,6 +110,10 @@ PK_API py_Callbacks* py_callbacks();
 
 /// Get the current source location of the frame.
 PK_API const char* py_Frame_sourceloc(py_Frame* frame, int* lineno);
+/// Python equivalent to `globals()` with respect to the given frame.
+PK_API void py_Frame_newglobals(py_Frame* frame, py_OutRef out);
+/// Python equivalent to `locals()` with respect to the given frame.
+PK_API void py_Frame_newlocals(py_Frame* frame, py_OutRef out);
 /// Get the function object of the frame.
 /// Returns `NULL` if not available.
 PK_API py_StackRef py_Frame_function(py_Frame* frame);
diff --git a/src/interpreter/frame.c b/src/interpreter/frame.c
index 260971c2..2e7f4824 100644
--- a/src/interpreter/frame.c
+++ b/src/interpreter/frame.c
@@ -192,6 +192,35 @@ const char* py_Frame_sourceloc(py_Frame* self, int* lineno) {
     return loc.src->filename->data;
 }
 
+void py_Frame_newglobals(py_Frame* frame, py_Ref out) {
+    if(!frame) {
+        pk_mappingproxy__namedict(out, &pk_current_vm->main);
+        return;
+    }
+    if(frame->globals->type == tp_module) {
+        pk_mappingproxy__namedict(out, frame->globals);
+    } else {
+        *out = *frame->globals;  // dict
+    }
+}
+
+void py_Frame_newlocals(py_Frame* frame, py_Ref out) {
+    if(!frame) {
+        py_newdict(out);
+        return;
+    }
+    if(frame->is_locals_special) {
+        switch(frame->locals->type) {
+            case tp_locals: frame = frame->locals->_ptr; break;
+            case tp_dict: *out = *frame->locals; return;
+            case tp_nil: py_newdict(out); return;
+            default: c11__unreachable();
+        }
+    }
+    FastLocals__to_dict(frame->locals, frame->co);
+    py_assign(out, py_retval());
+}
+
 py_StackRef py_Frame_function(py_Frame* self) {
     if(self->is_locals_special) return NULL;
     assert(self->p0->type == tp_function);
diff --git a/src/public/modules.c b/src/public/modules.c
index 51c5da68..3efb3627 100644
--- a/src/public/modules.c
+++ b/src/public/modules.c
@@ -501,33 +501,12 @@ static bool builtins_locals(int argc, py_Ref argv) {
 
 void py_newglobals(py_Ref out) {
     py_Frame* frame = pk_current_vm->top_frame;
-    if(!frame) {
-        pk_mappingproxy__namedict(out, &pk_current_vm->main);
-        return;
-    }
-    if(frame->globals->type == tp_module) {
-        pk_mappingproxy__namedict(out, frame->globals);
-    } else {
-        *out = *frame->globals;  // dict
-    }
+    py_Frame_newglobals(frame, out);
 }
 
 void py_newlocals(py_Ref out) {
     py_Frame* frame = pk_current_vm->top_frame;
-    if(!frame) {
-        py_newdict(out);
-        return;
-    }
-    if(frame->is_locals_special) {
-        switch(frame->locals->type) {
-            case tp_locals: frame = frame->locals->_ptr; break;
-            case tp_dict: *out = *frame->locals; return;
-            case tp_nil: py_newdict(out); return;
-            default: c11__unreachable();
-        }
-    }
-    FastLocals__to_dict(frame->locals, frame->co);
-    py_assign(out, py_retval());
+    py_Frame_newlocals(frame, out);
 }
 
 static void pk_push_special_locals() {