From 79cafcf32ccf5a5dea3c3c7427832c695ab3c107 Mon Sep 17 00:00:00 2001 From: blueloveTH Date: Thu, 28 Mar 2024 15:48:59 +0800 Subject: [PATCH] improve `io` --- src/io.cpp | 61 ++++++++++++++++++++++++++++++++++-------------- tests/70_file.py | 32 +++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 19 deletions(-) diff --git a/src/io.cpp b/src/io.cpp index c9602b7f..af77086e 100644 --- a/src/io.cpp +++ b/src/io.cpp @@ -12,12 +12,10 @@ namespace pkpy{ struct FileIO { PY_CLASS(FileIO, io, FileIO) - Str file; - Str mode; FILE* fp; + bool is_text; - bool is_text() const { return mode != "rb" && mode != "wb" && mode != "ab"; } - FileIO(VM* vm, std::string file, std::string mode); + FileIO(VM* vm, const Str& file, const Str& mode); void close(); static void _register(VM* vm, PyObject* mod, PyObject* type); }; @@ -62,27 +60,34 @@ void FileIO::_register(VM* vm, PyObject* mod, PyObject* type){ vm->bind_constructor<3>(type, [](VM* vm, ArgsView args){ Type cls = PK_OBJ_GET(Type, args[0]); return vm->heap.gcnew(cls, vm, - py_cast(vm, args[1]).str(), - py_cast(vm, args[2]).str()); + py_cast(vm, args[1]), + py_cast(vm, args[2])); }); - vm->bind_method<0>(type, "read", [](VM* vm, ArgsView args){ - FileIO& io = CAST(FileIO&, args[0]); - fseek(io.fp, 0, SEEK_END); - int buffer_size = ftell(io.fp); + vm->bind(type, "read(self, size=-1)", [](VM* vm, ArgsView args){ + FileIO& io = PK_OBJ_GET(FileIO, args[0]); + i64 size = CAST(i64, args[1]); + i64 buffer_size; + if(size < 0){ + long current = ftell(io.fp); + fseek(io.fp, 0, SEEK_END); + buffer_size = ftell(io.fp); + fseek(io.fp, current, SEEK_SET); + }else{ + buffer_size = size; + } unsigned char* buffer = new unsigned char[buffer_size]; - fseek(io.fp, 0, SEEK_SET); size_t actual_size = io_fread(buffer, 1, buffer_size, io.fp); PK_ASSERT(actual_size <= buffer_size); // in text mode, CR may be dropped, which may cause `actual_size < buffer_size` Bytes b(buffer, actual_size); - if(io.is_text()) return VAR(b.str()); + if(io.is_text) return VAR(b.str()); return VAR(std::move(b)); }); vm->bind_method<1>(type, "write", [](VM* vm, ArgsView args){ - FileIO& io = CAST(FileIO&, args[0]); - if(io.is_text()){ + FileIO& io = PK_OBJ_GET(FileIO, args[0]); + if(io.is_text){ Str& s = CAST(Str&, args[1]); fwrite(s.data, 1, s.length(), io.fp); }else{ @@ -92,14 +97,30 @@ void FileIO::_register(VM* vm, PyObject* mod, PyObject* type){ return vm->None; }); + vm->bind_method<0>(type, "tell", [](VM* vm, ArgsView args){ + FileIO& io = PK_OBJ_GET(FileIO, args[0]); + long pos = ftell(io.fp); + if(pos == -1) vm->IOError(strerror(errno)); + return VAR(pos); + }); + + vm->bind_method<2>(type, "seek", [](VM* vm, ArgsView args){ + FileIO& io = PK_OBJ_GET(FileIO, args[0]); + long offset = CAST(long, args[1]); + int whence = CAST(int, args[2]); + int ret = fseek(io.fp, offset, whence); + if(ret != 0) vm->IOError(strerror(errno)); + return vm->None; + }); + vm->bind_method<0>(type, "close", [](VM* vm, ArgsView args){ - FileIO& io = CAST(FileIO&, args[0]); + FileIO& io = PK_OBJ_GET(FileIO, args[0]); io.close(); return vm->None; }); vm->bind_method<0>(type, "__exit__", [](VM* vm, ArgsView args){ - FileIO& io = CAST(FileIO&, args[0]); + FileIO& io = PK_OBJ_GET(FileIO, args[0]); io.close(); return vm->None; }); @@ -107,7 +128,8 @@ void FileIO::_register(VM* vm, PyObject* mod, PyObject* type){ vm->bind_method<0>(type, "__enter__", PK_LAMBDA(args[0])); } -FileIO::FileIO(VM* vm, std::string file, std::string mode): file(file), mode(mode) { +FileIO::FileIO(VM* vm, const Str& file, const Str& mode){ + this->is_text = mode.sv().find("b") == std::string::npos; fp = io_fopen(file.c_str(), mode.c_str()); if(!fp) vm->IOError(strerror(errno)); } @@ -121,6 +143,11 @@ void FileIO::close(){ void add_module_io(VM* vm){ PyObject* mod = vm->new_module("io"); FileIO::register_class(vm, mod); + + mod->attr().set("SEEK_SET", VAR(SEEK_SET)); + mod->attr().set("SEEK_CUR", VAR(SEEK_CUR)); + mod->attr().set("SEEK_END", VAR(SEEK_END)); + vm->bind(vm->builtins, "open(path, mode='r')", [](VM* vm, ArgsView args){ PK_LOCAL_STATIC StrName m_io("io"); PK_LOCAL_STATIC StrName m_FileIO("FileIO"); diff --git a/tests/70_file.py b/tests/70_file.py index 2af68e45..15570180 100644 --- a/tests/70_file.py +++ b/tests/70_file.py @@ -13,6 +13,34 @@ a.close() with open('123.txt', 'rt') as f: assert f.read() == '123456' +with open('123.txt', 'rt') as f: + assert f.read(3) == '123' + assert f.tell() == 3 + assert f.read(3) == '456' + assert f.tell() == 6 + assert f.read(3) == '' # EOF + assert f.tell() == 6 + +with open('123.txt', 'rb') as f: + assert f.read(2) == b'12' + assert f.tell() == 2 + assert f.read(2) == b'34' + assert f.tell() == 4 + assert f.read(2) == b'56' + assert f.tell() == 6 + assert f.read(2) == b'' # EOF + assert f.tell() == 6 + +# test fseek +with open('123.txt', 'rt') as f: + f.seek(0, io.SEEK_END) + assert f.tell() == 6 + assert f.read() == '' + f.seek(3, io.SEEK_SET) + assert f.tell() == 3 + assert f.read() == '456' + assert f.tell() == 6 + with open('123.txt', 'a') as f: f.write('测试') @@ -29,13 +57,13 @@ with open('123.bin', 'wb') as f: f.write('123'.encode()) f.write('测试'.encode()) -def f(): +def f_(): with open('123.bin', 'rb') as f: b = f.read() assert isinstance(b, bytes) assert b == '123测试'.encode() -f() +f_() assert os.path.exists('123.bin') os.remove('123.bin')