From b69bef09125a3ae9d41dd69c2627b9530e4e705a Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Sun, 19 Mar 2023 15:05:55 +0800
Subject: [PATCH] allow `import` from file

---
 src/ceval.h | 19 ++++++++++++-------
 src/io.h    | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/ceval.h b/src/ceval.h
index d189acb0..2161117f 100644
--- a/src/ceval.h
+++ b/src/ceval.h
@@ -5,6 +5,8 @@
 
 namespace pkpy{
 
+Str _read_file_cwd(const Str& name, bool* ok);
+
 PyVar VM::run_frame(Frame* frame){
     while(frame->has_next_bytecode()){
         const Bytecode& byte = frame->next_bytecode();
@@ -298,18 +300,21 @@ PyVar VM::run_frame(Frame* frame){
             StrName name = frame->co->names[byte.arg].first;
             PyVar* ext_mod = _modules.try_get(name);
             if(ext_mod == nullptr){
+                Str source;
                 auto it2 = _lazy_modules.find(name);
                 if(it2 == _lazy_modules.end()){
-                    _error("ImportError", "module " + name.str().escape(true) + " not found");
+                    bool ok = false;
+                    source = _read_file_cwd(name.str() + ".py", &ok);
+                    if(!ok) _error("ImportError", "module " + name.str().escape(true) + " not found");
                 }else{
-                    const Str& source = it2->second;
-                    CodeObject_ code = compile(source, name.str(), EXEC_MODE);
-                    PyVar new_mod = new_module(name);
-                    _exec(code, new_mod);
-                    frame->push(new_mod);
+                    source = it2->second;
                     _lazy_modules.erase(it2);
-                    new_mod->attr()._try_perfect_rehash();
                 }
+                CodeObject_ code = compile(source, name.str(), EXEC_MODE);
+                PyVar new_mod = new_module(name);
+                _exec(code, new_mod);
+                frame->push(new_mod);
+                new_mod->attr()._try_perfect_rehash();
             }else{
                 frame->push(*ext_mod);
             }
diff --git a/src/io.h b/src/io.h
index e342d900..cbbef646 100644
--- a/src/io.h
+++ b/src/io.h
@@ -10,6 +10,20 @@
 
 namespace pkpy{
 
+Str _read_file_cwd(const Str& name, bool* ok){
+    std::filesystem::path path(name.c_str());
+    bool exists = std::filesystem::exists(path);
+    if(!exists){
+        *ok = false;
+        return Str();
+    }
+    std::ifstream ifs(path);
+    std::string buffer((std::istreambuf_iterator<char>(ifs)), (std::istreambuf_iterator<char>()));
+    ifs.close();
+    *ok = true;
+    return Str(std::move(buffer));
+}
+
 struct FileIO {
     PY_CLASS(FileIO, io, FileIO)
 
@@ -138,6 +152,12 @@ void add_module_os(VM* vm){
 namespace pkpy{
 void add_module_io(VM* vm){}
 void add_module_os(VM* vm){}
+
+Str _read_file_cwd(const Str& name, bool* ok){
+    *ok = false;
+    return Str();
+}
+
 } // namespace pkpy
 
 #endif
\ No newline at end of file