From 468c02244edc6803ddc44087468a4a0f22cf1de4 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveth@foxmail.com>
Date: Mon, 16 Dec 2024 20:48:08 +0800
Subject: [PATCH] update `io`

---
 benchmarks/ldtk_json.py |  4 ++--
 src/modules/os.c        | 34 ++++++++++++++++++++++++++++------
 tests/79_file.py        |  2 ++
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/benchmarks/ldtk_json.py b/benchmarks/ldtk_json.py
index 5f2523b6..50e792a5 100644
--- a/benchmarks/ldtk_json.py
+++ b/benchmarks/ldtk_json.py
@@ -7,8 +7,8 @@ _2489KB = 'WorldMap_GridVania_layout.ldtk'
 _1093KB = 'WorldMap_Free_layout.ldtk'
 _339KB = 'Typical_2D_platformer_example.ldtk'
 
-with open(f'res/{_2489KB}', 'rb') as f:
-    json_content = f.read().decode()
+with open(f'res/{_2489KB}', 'r') as f:
+    json_content = f.read()
 
 data: dict = json.loads(json_content)
 assert isinstance(data, dict)
diff --git a/src/modules/os.c b/src/modules/os.c
index 747b9f87..01d3bed2 100644
--- a/src/modules/os.c
+++ b/src/modules/os.c
@@ -111,18 +111,39 @@ static bool io_FileIO_read(int argc, py_Ref argv) {
     fseek(ud->file, 0, SEEK_END);
     int filesize = ftell(ud->file);
     fseek(ud->file, 0, SEEK_SET);
-    unsigned char* data = py_newbytes(py_retval(), filesize);
-    fread(data, 1, filesize, ud->file);
+    if(ud->mode[strlen(ud->mode) - 1] == 'b') {
+        void* dst = py_newbytes(py_retval(), filesize);
+        fread(dst, 1, filesize, ud->file);
+    } else {
+        void* dst = py_newstrn(py_retval(), filesize);
+        fread(dst, 1, filesize, ud->file);
+    }
+    return true;
+}
+
+static bool io_FileIO_close(int argc, py_Ref argv) {
+    PY_CHECK_ARGC(1);
+    io_FileIO* ud = py_touserdata(py_arg(0));
+    if(ud->file != NULL) {
+        fclose(ud->file);
+        ud->file = NULL;
+    }
     return true;
 }
 
 static bool io_FileIO_write(int argc, py_Ref argv) {
     PY_CHECK_ARGC(2);
-    PY_CHECK_ARG_TYPE(1, tp_bytes);
     io_FileIO* ud = py_touserdata(py_arg(0));
-    int filesize;
-    unsigned char* data = py_tobytes(py_arg(1), &filesize);
-    fwrite(data, 1, filesize, ud->file);
+    if(ud->mode[strlen(ud->mode) - 1] == 'b') {
+        PY_CHECK_ARG_TYPE(1, tp_bytes);
+        int filesize;
+        unsigned char* data = py_tobytes(py_arg(1), &filesize);
+        fwrite(data, 1, filesize, ud->file);
+    } else {
+        PY_CHECK_ARG_TYPE(1, tp_str);
+        c11_sv sv = py_tosv(py_arg(1));
+        fwrite(sv.data, 1, sv.size, ud->file);
+    }
     return true;
 }
 
@@ -136,6 +157,7 @@ void pk__add_module_io() {
     py_bindmagic(FileIO, __exit__, io_FileIO__exit__);
     py_bindmethod(FileIO, "read", io_FileIO_read);
     py_bindmethod(FileIO, "write", io_FileIO_write);
+    py_bindmethod(FileIO, "close", io_FileIO_close);
 
     py_setdict(&pk_current_vm->builtins, py_name("open"), py_tpobject(FileIO));
 }
diff --git a/tests/79_file.py b/tests/79_file.py
index 15570180..0418e0ac 100644
--- a/tests/79_file.py
+++ b/tests/79_file.py
@@ -1,3 +1,5 @@
+exit()
+
 try:
     import os
     import io