From a17755ee8ba2982505d349e0f06acdabe943d33e Mon Sep 17 00:00:00 2001
From: Dave Hylands <dhylands@gmail.com>
Date: Wed, 23 Dec 2015 19:11:27 -0800
Subject: [PATCH] py: Add ustruct.pack_into and unpack_from
---
py/modstruct.c | 94 ++++++++++++++++++++++++++++++++---------
py/qstrdefs.h | 2 +
tests/basics/struct1.py | 41 ++++++++++++++++++
3 files changed, 118 insertions(+), 19 deletions(-)
diff --git a/py/modstruct.c b/py/modstruct.c
index 727a7ab00..cd3209738 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 7e2109006..d27525f68 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 61c08b74e..1afea5a96 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('<bbb', buf, 3, 0x41, 0x42, 0x43)
+print(buf)
+struct.pack_into('<bbb', buf, -6, 0x44, 0x45, 0x46)
+print(buf)
+
+try:
+ struct.pack_into('<bbb', buf, 7, 0x41, 0x42, 0x43)
+except:
+ print('struct.error')
+try:
+ struct.pack_into('<bbb', buf, -10, 0x41, 0x42, 0x43)
+except:
+ print('struct.error')
+
+# unpack_from
+buf = b'0123456789'
+print(struct.unpack_from('<b', buf, 4))
+print(struct.unpack_from('<b', buf, -4))
+try:
+ print(struct.unpack_from('<b', buf, 10))
+except:
+ print('struct.error')
+try:
+ print(struct.unpack_from('<b', buf, -11))
+except:
+ print('struct.error')
+
+# pack with too many args, not checked by uPy
+#try:
+# print(struct.pack('ii', 1, 2, 3))
+#except:
+# print('struct.error')
+
+# pack with too few args, not checked by uPy
+#try:
+# print(struct.pack('ii', 1))
+#except:
+# print('struct.error')
--
GitLab