From aaa6d60404fea0e3c950e9121b391bb60cbfc45e Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 10 Sep 2023 03:22:01 +0800
Subject: [PATCH] add `__module__` for types and fix pickle bugs

---
 include/pocketpy/vm.h |  1 +
 python/pickle.py      | 22 ++++++++++++++--------
 src/pocketpy.cpp      |  5 +++++
 src/vm.cpp            |  7 ++++---
 tests/81_pickle.py    |  4 +++-
 5 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/include/pocketpy/vm.h b/include/pocketpy/vm.h
index 7e375cd0..9eedea93 100644
--- a/include/pocketpy/vm.h
+++ b/include/pocketpy/vm.h
@@ -51,6 +51,7 @@ typedef PyObject* (*BinaryFuncC)(VM*, PyObject*, PyObject*);
 struct PyTypeInfo{
     PyObject* obj;
     Type base;
+    PyObject* mod;
     Str name;
     bool subclass_enabled;
 
diff --git a/python/pickle.py b/python/pickle.py
index 74e20bf5..bc23785f 100644
--- a/python/pickle.py
+++ b/python/pickle.py
@@ -2,15 +2,13 @@ import json
 import builtins
 
 _BASIC_TYPES = [int, float, str, bool, type(None)]
+_MOD_T_SEP = "@"
 
 def _find_class(path: str):
-    if "." not in path:
-        g = globals()
-        if path in g:
-            return g[path]
+    if _MOD_T_SEP not in path:
         return builtins.__dict__[path]
-    modname, name = path.split(".")
-    return __import__(modname).__dict__[name]
+    modpath, name = path.split(_MOD_T_SEP)
+    return __import__(modpath).__dict__[name]
 
 def _find__new__(cls):
     while cls is not None:
@@ -26,11 +24,19 @@ class _Pickler:
         self.raw_memo = {}  # id -> int
         self.memo = []      # int -> object
 
+    def _type_id(self, o: type):
+        assert type(o) is type
+        name = o.__name__
+        mod = o.__module__
+        if mod is not None:
+            name = mod.__path__ + _MOD_T_SEP + name
+        return name
+
     def wrap(self, o):
         if type(o) in _BASIC_TYPES:
             return o
         if type(o) is type:
-            return ["type", o.__name__]
+            return ["type", self._type_id(o)]
 
         index = self.raw_memo.get(id(o), None)
         if index is not None:
@@ -59,7 +65,7 @@ class _Pickler:
             ret.append([[self.wrap(k), self.wrap(v)] for k,v in o.items()])
             return [index]
         
-        _0 = o.__class__.__name__
+        _0 = self._type_id(type(o))
 
         if hasattr(o, "__getnewargs__"):
             _1 = o.__getnewargs__()     # an iterable
diff --git a/src/pocketpy.cpp b/src/pocketpy.cpp
index d3870a6f..36315021 100644
--- a/src/pocketpy.cpp
+++ b/src/pocketpy.cpp
@@ -1540,6 +1540,11 @@ void VM::post_init(){
         const PyTypeInfo& info = vm->_all_types[PK_OBJ_GET(Type, args[0])];
         return VAR(info.name);
     });
+    bind_property(_t(tp_type), "__module__", [](VM* vm, ArgsView args){
+        const PyTypeInfo& info = vm->_all_types[PK_OBJ_GET(Type, args[0])];
+        if(info.mod == nullptr) return vm->None;
+        return info.mod;
+    });
     bind_property(_t(tp_bound_method), "__self__", [](VM* vm, ArgsView args){
         return CAST(BoundMethod&, args[0]).self;
     });
diff --git a/src/vm.cpp b/src/vm.cpp
index 14b3ac94..e44d2cdf 100644
--- a/src/vm.cpp
+++ b/src/vm.cpp
@@ -131,7 +131,8 @@ namespace pkpy{
         PyTypeInfo info{
             obj,
             base,
-            (mod!=nullptr && mod!=builtins) ? Str(OBJ_NAME(mod)+"."+name.sv()): name.sv(),
+            mod,
+            name.sv(),
             subclass_enabled,
         };
         if(mod != nullptr) mod->attr().set(name, obj);
@@ -636,8 +637,8 @@ void VM::_log_s_data(const char* title) {
 #endif
 
 void VM::init_builtin_types(){
-    _all_types.push_back({heap._new<Type>(Type(1), Type(0)), -1, "object", true});
-    _all_types.push_back({heap._new<Type>(Type(1), Type(1)), 0, "type", false});
+    _all_types.push_back({heap._new<Type>(Type(1), Type(0)), -1, nullptr, "object", true});
+    _all_types.push_back({heap._new<Type>(Type(1), Type(1)), 0, nullptr, "type", false});
     tp_object = 0; tp_type = 1;
 
     tp_int = _new_type_object("int");
diff --git a/tests/81_pickle.py b/tests/81_pickle.py
index bd922dca..392fb84e 100644
--- a/tests/81_pickle.py
+++ b/tests/81_pickle.py
@@ -1,7 +1,9 @@
 from pickle import dumps, loads, _wrap, _unwrap
 
 def test(x):
-    ok = x == loads(dumps(x))
+    y = dumps(x)
+    # print(y.decode())
+    ok = x == loads(y)
     if not ok:
         _0 = _wrap(x)
         _1 = _unwrap(_0)