From 4a8776ec6c2b8ff7bf2179d320a0e5fbeefe4fa4 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Sun, 24 Nov 2024 12:43:08 +0800
Subject: [PATCH] add `inspect` module

---
 include/pocketpy/interpreter/modules.h |  2 ++
 include/pocketpy/pocketpy.h            |  2 +-
 include/typings/inspect.pyi            |  1 +
 src/interpreter/vm.c                   |  1 +
 src/modules/gc.c                       |  1 +
 src/modules/inspect.c                  | 28 +++++++++++++++
 tests/97_inspect.py                    | 49 ++++++++++++++++++++++++++
 7 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 include/typings/inspect.pyi
 create mode 100644 src/modules/inspect.c
 create mode 100644 tests/97_inspect.py

diff --git a/include/pocketpy/interpreter/modules.h b/include/pocketpy/interpreter/modules.h
index 1b86291b..178a5df0 100644
--- a/include/pocketpy/interpreter/modules.h
+++ b/include/pocketpy/interpreter/modules.h
@@ -11,6 +11,8 @@ void pk__add_module_time();
 void pk__add_module_easing();
 void pk__add_module_traceback();
 void pk__add_module_enum();
+void pk__add_module_inspect();
+
 void pk__add_module_linalg();
 void pk__add_module_array2d();
 
diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h
index fa2248bf..02bdf21f 100644
--- a/include/pocketpy/pocketpy.h
+++ b/include/pocketpy/pocketpy.h
@@ -672,7 +672,7 @@ enum py_PredefinedTypes {
     tp_module,
     tp_function,
     tp_nativefunc,
-    tp_boundmethod,
+    tp_boundmethod,    // 2 slots (self, func)
     tp_super,          // 1 slot + py_Type
     tp_BaseException,  // 2 slots (arg + inner_exc)
     tp_Exception,
diff --git a/include/typings/inspect.pyi b/include/typings/inspect.pyi
new file mode 100644
index 00000000..93c18d98
--- /dev/null
+++ b/include/typings/inspect.pyi
@@ -0,0 +1 @@
+def isgeneratorfunction(obj) -> bool: ...
\ No newline at end of file
diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c
index f78bb402..3d948edc 100644
--- a/src/interpreter/vm.c
+++ b/src/interpreter/vm.c
@@ -215,6 +215,7 @@ void VM__ctor(VM* self) {
     pk__add_module_easing();
     pk__add_module_traceback();
     pk__add_module_enum();
+    pk__add_module_inspect();
 
     pk__add_module_conio();
     pk__add_module_pkpy();
diff --git a/src/modules/gc.c b/src/modules/gc.c
index f00c7507..50527a88 100644
--- a/src/modules/gc.c
+++ b/src/modules/gc.c
@@ -6,6 +6,7 @@
 #include "pocketpy/interpreter/vm.h"
 
 static bool gc_collect(int argc, py_Ref argv){
+    PY_CHECK_ARGC(0);
     ManagedHeap* heap = &pk_current_vm->heap;
     int res = ManagedHeap__collect(heap);
     py_newint(py_retval(), res);
diff --git a/src/modules/inspect.c b/src/modules/inspect.c
new file mode 100644
index 00000000..d32a4740
--- /dev/null
+++ b/src/modules/inspect.c
@@ -0,0 +1,28 @@
+#include "pocketpy/pocketpy.h"
+
+#include "pocketpy/common/utils.h"
+#include "pocketpy/objects/object.h"
+#include "pocketpy/common/sstream.h"
+#include "pocketpy/interpreter/vm.h"
+
+static bool inspect_isgeneratorfunction(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    py_Ref obj = argv;
+    if(py_istype(argv, tp_boundmethod)) {
+        py_TValue* slots = PyObject__slots(argv->_obj);
+        obj = &slots[1];  // callable
+    }
+    if(py_istype(obj, tp_function)) {
+        Function* fn = py_touserdata(obj);
+        py_newbool(py_retval(), fn->decl->type == FuncType_GENERATOR);
+    } else {
+        py_newbool(py_retval(), false);
+    }
+    return true;
+}
+
+void pk__add_module_inspect() {
+    py_Ref mod = py_newmodule("inspect");
+
+    py_bindfunc(mod, "isgeneratorfunction", inspect_isgeneratorfunction);
+}
\ No newline at end of file
diff --git a/tests/97_inspect.py b/tests/97_inspect.py
new file mode 100644
index 00000000..6ce9a108
--- /dev/null
+++ b/tests/97_inspect.py
@@ -0,0 +1,49 @@
+from inspect import isgeneratorfunction
+
+def f(a, b):
+    return a + b
+
+assert not isgeneratorfunction(f)
+
+def g(a, b):
+    yield a
+    yield b
+
+assert isgeneratorfunction(g)
+
+class A:
+    @staticmethod
+    def non_gen(a, b):
+        return a + b
+    @staticmethod
+    def gen(a, b):
+        yield a
+        yield b
+    @classmethod
+    def non_gen_class(cls, a, b):
+        return a + b
+    @classmethod
+    def gen_class(cls, a, b):
+        yield a
+        yield b
+    def not_gen_instance(self, a, b):
+        return a + b
+    def gen_instance(self, a, b):
+        yield a
+        yield b
+
+a = A()
+assert not isgeneratorfunction(a.non_gen)
+assert isgeneratorfunction(a.gen)
+assert not isgeneratorfunction(A.non_gen)
+assert isgeneratorfunction(A.gen)
+
+assert not isgeneratorfunction(a.non_gen_class)
+assert isgeneratorfunction(a.gen_class)
+assert not isgeneratorfunction(A.non_gen_class)
+assert isgeneratorfunction(A.gen_class)
+
+assert not isgeneratorfunction(a.not_gen_instance)
+assert isgeneratorfunction(a.gen_instance)
+assert not isgeneratorfunction(A.not_gen_instance)
+assert isgeneratorfunction(A.gen_instance)