support struct.unpack()

pikasTech · Aug 17, 2023 · 4665de2 · 4665de2
1 parent cf7990a
commit 4665de2
Show file tree

Hide file tree

Showing 10 changed files with 122 additions and 9 deletions.
diff --git a/examples/struct/unpack.py b/examples/struct/unpack.py
@@ -0,0 +1,27 @@
+import struct
+unpacked_bytes = struct.unpack('bbhhh', b'\x01\x02\x03\x04\x05\x06\x07\x08')
+assert unpacked_bytes[0] == 1
+assert unpacked_bytes[1] == 2
+assert unpacked_bytes[2] == 1027
+assert unpacked_bytes[3] == 1541
+assert unpacked_bytes[4] == 2055
+
+unpacked_bytes = struct.unpack('If', b'\x01\x00\x00\x00\x42\x28\xAE\x47')
+
+assert unpacked_bytes[0] == 1
+assert unpacked_bytes[1] == 89168.515625
+
+unpacked_string_numbers = struct.unpack(
+    '5sIf', b'Hello\x00\x00\x00\x01\x00\x00\x00B(\xaeG')
+assert unpacked_string_numbers[0] == b'Hello'
+assert unpacked_string_numbers[1] == 1
+assert unpacked_string_numbers[2] == 89168.515625
+
+unpacked_mixed_data = struct.unpack(
+    'BffI', b'\x01\x00\x00\x00+R\x9aD\x00\x00\x80?P\\\x01\x00')
+assert unpacked_mixed_data[0] == 1
+assert unpacked_mixed_data[1] == 1234.5677490234375
+assert unpacked_mixed_data[2] == 1
+assert unpacked_mixed_data[3] == 89168
+
+print('PASS')
diff --git a/package/struct/_struct.c b/package/struct/_struct.c
@@ -113,7 +113,7 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
     return MP_OBJ_NEW_SMALL_INT(size);
 }
 
-STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
+STATIC mp_obj_t struct_unpack_from(size_t n_args, mp_obj_t* args) {
     // unpack requires that the buffer be exactly the right size.
     // unpack_from requires that the buffer be "big enough".
     // Since we implement unpack and unpack_from using the same function
@@ -122,7 +122,8 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
     size_t total_sz;
     size_t num_items = calc_size_items(fmt, &total_sz);
     char fmt_type = get_fmt_type(&fmt);
-    mp_obj_tuple_t* res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL));
+    mp_obj_t tuple = mp_obj_new_tuple(num_items, NULL);
+    mp_obj_tuple_t* res = MP_OBJ_TO_PTR(tuple);
     mp_buffer_info_t bufinfo;
     mp_get_buffer_raise((Arg*)args[1], &bufinfo, MP_BUFFER_READ);
     byte* p = bufinfo.buf;
@@ -166,7 +167,24 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
         }
         fmt++;
     }
-    return MP_OBJ_FROM_PTR(res);
+    mp_obj_t ret = MP_OBJ_FROM_TUPLE(res);
+    mp_buff_info_deinit(&bufinfo);
+    mp_obj_tuple_deinit(tuple);
+    return ret;
+}
+
+Arg* _struct_unpack(PikaObj* self, char* fmt, Arg* data, int offset) {
+    Arg* aFmt = arg_newStr(fmt);
+    Arg* aOffset = arg_newInt(offset);
+    Arg** args = pikaMalloc(sizeof(Arg) * 3);
+    args[0] = aFmt;
+    args[1] = data;
+    args[2] = aOffset;
+    Arg* tuple = struct_unpack_from(3, args);
+    arg_deinit(aFmt);
+    arg_deinit(aOffset);
+    pikaFree(args, sizeof(Arg) * 3);
+    return tuple;
 }
 
 // This function assumes there is enough room in p to store all the values

diff --git a/package/struct/_struct.pyi b/package/struct/_struct.pyi
@@ -1 +1,2 @@
 def pack(fmt: str, *vars) -> bytes: ...
+def unpack(fmt: str, data: any, offset: int) -> any: ...
diff --git a/package/struct/struct.py b/package/struct/struct.py
@@ -1,4 +1,13 @@
 import _struct
 
-def pack(fmt, *args):
+
+def pack(fmt: str, *args) -> bytes:
     return _struct.pack(fmt, *args)
+
+
+def unpack(fmt: str, data: bytes, offset=0) -> tuple:
+    return _struct.unpack(fmt, data, offset)
+
+
+def unpack_from(fmt: str, data: bytes, offset=0) -> tuple:
+    return _struct.unpack(fmt, data, offset)
diff --git a/port/linux/.vscode/launch.json b/port/linux/.vscode/launch.json
@@ -12,7 +12,7 @@
             // "program": "${workspaceFolder}/build/boot/demo06-pikamain/pikascript_demo06-pikamain",
             "args": [
                 // "--gtest_filter=vm.keyword_2"
-                "--gtest_filter=struct.pack"
+                "--gtest_filter=struct.unpack"
             ],
             "stopAtEntry": false,
             "cwd": "${workspaceFolder}",

diff --git a/port/linux/package/pikascript/_struct.pyi b/port/linux/package/pikascript/_struct.pyi
@@ -1 +1,2 @@
 def pack(fmt: str, *vars) -> bytes: ...
+def unpack(fmt: str, data: any, offset: int) -> any: ...
diff --git a/port/linux/package/pikascript/pikascript-lib/struct/_struct.c b/port/linux/package/pikascript/pikascript-lib/struct/_struct.c
@@ -113,7 +113,7 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
     return MP_OBJ_NEW_SMALL_INT(size);
 }
 
-STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
+STATIC mp_obj_t struct_unpack_from(size_t n_args, mp_obj_t* args) {
     // unpack requires that the buffer be exactly the right size.
     // unpack_from requires that the buffer be "big enough".
     // Since we implement unpack and unpack_from using the same function
@@ -122,7 +122,8 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
     size_t total_sz;
     size_t num_items = calc_size_items(fmt, &total_sz);
     char fmt_type = get_fmt_type(&fmt);
-    mp_obj_tuple_t* res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL));
+    mp_obj_t tuple = mp_obj_new_tuple(num_items, NULL);
+    mp_obj_tuple_t* res = MP_OBJ_TO_PTR(tuple);
     mp_buffer_info_t bufinfo;
     mp_get_buffer_raise((Arg*)args[1], &bufinfo, MP_BUFFER_READ);
     byte* p = bufinfo.buf;
@@ -166,7 +167,24 @@ STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t* args) {
         }
         fmt++;
     }
-    return MP_OBJ_FROM_PTR(res);
+    mp_obj_t ret = MP_OBJ_FROM_TUPLE(res);
+    mp_buff_info_deinit(&bufinfo);
+    mp_obj_tuple_deinit(tuple);
+    return ret;
+}
+
+Arg* _struct_unpack(PikaObj* self, char* fmt, Arg* data, int offset) {
+    Arg* aFmt = arg_newStr(fmt);
+    Arg* aOffset = arg_newInt(offset);
+    Arg** args = pikaMalloc(sizeof(Arg) * 3);
+    args[0] = aFmt;
+    args[1] = data;
+    args[2] = aOffset;
+    Arg* tuple = struct_unpack_from(3, args);
+    arg_deinit(aFmt);
+    arg_deinit(aOffset);
+    pikaFree(args, sizeof(Arg) * 3);
+    return tuple;
 }
 
 // This function assumes there is enough room in p to store all the values

diff --git a/port/linux/package/pikascript/struct.py b/port/linux/package/pikascript/struct.py
@@ -1,4 +1,13 @@
 import _struct
 
-def pack(fmt, *args):
+
+def pack(fmt: str, *args) -> bytes:
     return _struct.pack(fmt, *args)
+
+
+def unpack(fmt: str, data: bytes, offset=0) -> tuple:
+    return _struct.unpack(fmt, data, offset)
+
+
+def unpack_from(fmt: str, data: bytes, offset=0) -> tuple:
+    return _struct.unpack(fmt, data, offset)
diff --git a/port/linux/test/module-test.cpp b/port/linux/test/module-test.cpp
@@ -629,7 +629,10 @@ TEST_RUN_SINGLE_FILE(modbus,
 
 TEST_RUN_SINGLE_FILE_PASS(proxy, proxy3, "test/python/proxy/proxy3.py")
 
+#if PIKA_FLOAT_TYPE_DOUBLE
 TEST_RUN_SINGLE_FILE_PASS(struct, pack, "test/python/struct/pack.py")
+TEST_RUN_SINGLE_FILE_PASS(struct, unpack, "test/python/struct/unpack.py")
+#endif
 
 #endif
 

diff --git a/port/linux/test/python/struct/unpack.py b/port/linux/test/python/struct/unpack.py
@@ -0,0 +1,27 @@
+import struct
+unpacked_bytes = struct.unpack('bbhhh', b'\x01\x02\x03\x04\x05\x06\x07\x08')
+assert unpacked_bytes[0] == 1
+assert unpacked_bytes[1] == 2
+assert unpacked_bytes[2] == 1027
+assert unpacked_bytes[3] == 1541
+assert unpacked_bytes[4] == 2055
+
+unpacked_bytes = struct.unpack('If', b'\x01\x00\x00\x00\x42\x28\xAE\x47')
+
+assert unpacked_bytes[0] == 1
+assert unpacked_bytes[1] == 89168.515625
+
+unpacked_string_numbers = struct.unpack(
+    '5sIf', b'Hello\x00\x00\x00\x01\x00\x00\x00B(\xaeG')
+assert unpacked_string_numbers[0] == b'Hello'
+assert unpacked_string_numbers[1] == 1
+assert unpacked_string_numbers[2] == 89168.515625
+
+unpacked_mixed_data = struct.unpack(
+    'BffI', b'\x01\x00\x00\x00+R\x9aD\x00\x00\x80?P\\\x01\x00')
+assert unpacked_mixed_data[0] == 1
+assert unpacked_mixed_data[1] == 1234.5677490234375
+assert unpacked_mixed_data[2] == 1
+assert unpacked_mixed_data[3] == 89168
+
+print('PASS')