py: Add ustruct.pack_into and unpack_from

1 files changed, 75 insertions, 19 deletions

diff --git a/py/modstruct.c b/py/modstruct.c
index 727a7ab006..cd32097388 100644
--- a/py/modstruct.c
+++ b/py/modstruct.c
@@ -132,22 +132,45 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
 }
 MP_DEFINE_CONST_FUN_OBJ_1(struct_calcsize_obj, struct_calcsize);
 
-STATIC mp_obj_t struct_unpack(mp_obj_t fmt_in, mp_obj_t data_in) {
-    // TODO: "The buffer must contain exactly the amount of data required by the format (len(bytes) must equal calcsize(fmt))."
-    const char *fmt = mp_obj_str_get_str(fmt_in);
+STATIC mp_obj_t struct_unpack_from(size_t n_args, const 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
+    // we relax the "exact" requirement, and only implement "big enough".
+    const char *fmt = mp_obj_str_get_str(args[0]);
     char fmt_type = get_fmt_type(&fmt);
-    uint size = calcsize_items(fmt);
-    mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(size, NULL));
+    uint num_items = calcsize_items(fmt);
+    mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL));
     mp_buffer_info_t bufinfo;
-    mp_get_buffer_raise(data_in, &bufinfo, MP_BUFFER_READ);
+    mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ);
     byte *p = bufinfo.buf;
+    byte *end_p = &p[bufinfo.len];
+    mp_int_t offset = 0;
 
-    for (uint i = 0; i < size;) {
+    if (n_args > 2) {
+        // offset arg provided
+        offset = mp_obj_get_int(args[2]);
+        if (offset < 0) {
+            // negative offsets are relative to the end of the buffer
+            offset = bufinfo.len + offset;
+            if (offset < 0) {
+                nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
+            }
+        }
+        p += offset;
+    }
+
+    for (uint i = 0; i < num_items;) {
+        if (*fmt == '\0') {
+            break;
+        }
         mp_uint_t sz = 1;
         if (unichar_isdigit(*fmt)) {
             sz = get_fmt_num(&fmt);
         }
-
+        if (p + sz > end_p) {
+            nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
+        }
         mp_obj_t item;
         if (*fmt == 's') {
             item = mp_obj_new_bytes(p, sz);
@@ -163,23 +186,24 @@ STATIC mp_obj_t struct_unpack(mp_obj_t fmt_in, mp_obj_t data_in) {
     }
     return MP_OBJ_FROM_PTR(res);
 }
-MP_DEFINE_CONST_FUN_OBJ_2(struct_unpack_obj, struct_unpack);
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_unpack_from_obj, 2, 3, struct_unpack_from);
 
-STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
-    // TODO: "The arguments must match the values required by the format exactly."
-    const char *fmt = mp_obj_str_get_str(args[0]);
+STATIC void struct_pack_into_internal(mp_obj_t fmt_in, byte *p, byte* end_p, size_t n_args, const mp_obj_t *args) {
+    const char *fmt = mp_obj_str_get_str(fmt_in);
     char fmt_type = get_fmt_type(&fmt);
-    mp_int_t size = MP_OBJ_SMALL_INT_VALUE(struct_calcsize(args[0]));
-    vstr_t vstr;
-    vstr_init_len(&vstr, size);
-    byte *p = (byte*)vstr.buf;
-    memset(p, 0, size);
 
-    for (mp_uint_t i = 1; i < n_args;) {
+    size_t i;
+    for (i = 0; i < n_args;) {
         mp_uint_t sz = 1;
+        if (*fmt == '\0') {
+            break;
+        }
         if (unichar_isdigit(*fmt)) {
             sz = get_fmt_num(&fmt);
         }
+        if (p + sz > end_p) {
+            nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
+        }
 
         if (*fmt == 's') {
             mp_buffer_info_t bufinfo;
@@ -198,16 +222,48 @@ STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
         }
         fmt++;
     }
+}
 
+STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
+    // TODO: "The arguments must match the values required by the format exactly."
+    mp_int_t size = MP_OBJ_SMALL_INT_VALUE(struct_calcsize(args[0]));
+    vstr_t vstr;
+    vstr_init_len(&vstr, size);
+    byte *p = (byte*)vstr.buf;
+    memset(p, 0, size);
+    byte *end_p = &p[size];
+    struct_pack_into_internal(args[0], p, end_p, n_args - 1, &args[1]);
     return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
 }
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_obj, 1, MP_OBJ_FUN_ARGS_MAX, struct_pack);
 
+STATIC mp_obj_t struct_pack_into(size_t n_args, const mp_obj_t *args) {
+    mp_buffer_info_t bufinfo;
+    mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_WRITE);
+    mp_int_t offset = mp_obj_get_int(args[2]);
+    if (offset < 0) {
+        // negative offsets are relative to the end of the buffer
+        offset = (mp_int_t)bufinfo.len + offset;
+        if (offset < 0) {
+            nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
+        }
+    }
+    byte *p = (byte *)bufinfo.buf;
+    byte *end_p = &p[bufinfo.len];
+    p += offset;
+
+    struct_pack_into_internal(args[0], p, end_p, n_args - 3, &args[3]);
+    return mp_const_none;
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_into_obj, 3, MP_OBJ_FUN_ARGS_MAX, struct_pack_into);
+
 STATIC const mp_rom_map_elem_t mp_module_struct_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_ustruct) },
     { MP_ROM_QSTR(MP_QSTR_calcsize), MP_ROM_PTR(&struct_calcsize_obj) },
     { MP_ROM_QSTR(MP_QSTR_pack), MP_ROM_PTR(&struct_pack_obj) },
-    { MP_ROM_QSTR(MP_QSTR_unpack), MP_ROM_PTR(&struct_unpack_obj) },
+    { MP_ROM_QSTR(MP_QSTR_pack_into), MP_ROM_PTR(&struct_pack_into_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unpack), MP_ROM_PTR(&struct_unpack_from_obj) },
+    { MP_ROM_QSTR(MP_QSTR_unpack_from), MP_ROM_PTR(&struct_unpack_from_obj) },
 };
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_struct_globals, mp_module_struct_globals_table);