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); diff --git a/py/qstrdefs.h b/py/qstrdefs.h index 7e21090063..d27525f685 100644 --- a/py/qstrdefs.h +++ b/py/qstrdefs.h @@ -491,7 +491,9 @@ Q(print_exception) Q(struct) Q(ustruct) Q(pack) +Q(pack_into) Q(unpack) +Q(unpack_from) Q(calcsize) #endif diff --git a/tests/basics/struct1.py b/tests/basics/struct1.py index 61c08b74e4..1afea5a96f 100644 --- a/tests/basics/struct1.py +++ b/tests/basics/struct1.py @@ -66,3 +66,44 @@ except TypeError: # but later were implemented for all. print(struct.unpack("<3B2h", b"foo\x12\x34\xff\xff")) print(struct.pack("<3B", 1, 2, 3)) + +# pack_into +buf = bytearray(b'>>>123<<<') +struct.pack_into('