Improve struct compatibility with CPython

1 files changed, 40 insertions, 23 deletions

diff --git a/shared-bindings/struct/__init__.c b/shared-bindings/struct/__init__.c
index 093597785..9240a15bb 100644
--- a/shared-bindings/struct/__init__.c
+++ b/shared-bindings/struct/__init__.c
@@ -51,7 +51,7 @@
 //|
 //| Supported size/byte order prefixes: *@*, *<*, *>*, *!*.
 //|
-//| Supported format codes: *b*, *B*, *h*, *H*, *i*, *I*, *l*, *L*, *q*, *Q*,
+//| Supported format codes: *b*, *B*, *x*, *h*, *H*, *i*, *I*, *l*, *L*, *q*, *Q*,
 //| *s*, *P*, *f*, *d* (the latter 2 depending on the floating-point support).
 //|
 
@@ -74,7 +74,6 @@ MP_DEFINE_CONST_FUN_OBJ_1(struct_calcsize_obj, struct_calcsize);
 //|
 
 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);
@@ -115,49 +114,67 @@ MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_into_obj, 3, MP_OBJ_FUN_ARGS_MAX
 //| .. function:: unpack(fmt, data)
 //|
 //|   Unpack from the data according to the format string fmt. The return value
-//|   is a tuple of the unpacked values.
+//|   is a tuple of the unpacked values. The buffer size must match the size
+//|   required by the format.
 //|
 
-//| .. function:: unpack_from(fmt, data, offset)
+STATIC mp_obj_t struct_unpack(size_t n_args, const mp_obj_t *args) {
+    mp_buffer_info_t bufinfo;
+    mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ);
+    byte *p = bufinfo.buf;
+    byte *end_p = &p[bufinfo.len];
+
+    // true means check the size must be exactly right.
+    return MP_OBJ_FROM_PTR(shared_modules_struct_unpack_from(args[0] , p, end_p, true));
+}
+MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_unpack_obj, 2, 3, struct_unpack);
+
+//| .. function:: unpack_from(fmt, data, offset=0)
 //|
 //|   Unpack from the data starting at offset according to the format string fmt.
 //|   offset may be negative to count from the end of buffer. The return value is
-//|   a tuple of the unpacked values.
+//|   a tuple of the unpacked values. The buffer size must be at least as big
+//|   as the size required by the form.
 //|
 
-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".
+STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    enum { ARG_format, ARG_buffer, ARG_offset };
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_format, MP_ARG_REQUIRED | MP_ARG_OBJ },
+        { MP_QSTR_buffer, MP_ARG_REQUIRED | MP_ARG_OBJ },
+        { MP_QSTR_offset, MP_ARG_INT, {.u_int = 0} },
+    };
+
+    mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
+    mp_arg_parse_all(n_args, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
+
     mp_buffer_info_t bufinfo;
-    mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ);
+    mp_get_buffer_raise(args[ARG_buffer].u_obj, &bufinfo, MP_BUFFER_READ);
     byte *p = bufinfo.buf;
     byte *end_p = &p[bufinfo.len];
 
-    if (n_args > 2) {
-        mp_int_t offset = mp_obj_get_int(args[2]);
-        // offset arg provided
+    mp_int_t offset = args[ARG_offset].u_int;
+    if (offset < 0) {
+        // negative offsets are relative to the end of the buffer
+        offset = bufinfo.len + offset;
         if (offset < 0) {
-            // negative offsets are relative to the end of the buffer
-            offset = bufinfo.len + offset;
-            if (offset < 0) {
-                mp_raise_RuntimeError(translate("buffer too small"));
-            }
+            mp_raise_RuntimeError(translate("buffer too small"));
         }
-        p += offset;
     }
+    p += offset;
 
-    return MP_OBJ_FROM_PTR(shared_modules_struct_unpack_from(args[0] , p, end_p));
+    // false means the size doesn't have to be exact. struct.unpack_from() only requires
+    // that be buffer be big enough.
+    return MP_OBJ_FROM_PTR(shared_modules_struct_unpack_from(args[ARG_format].u_obj, p, end_p, false));
 }
-MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_unpack_from_obj, 2, 3, struct_unpack_from);
+MP_DEFINE_CONST_FUN_OBJ_KW(struct_unpack_from_obj, 0, struct_unpack_from);
 
 STATIC const mp_rom_map_elem_t mp_module_struct_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_struct) },
     { 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_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), MP_ROM_PTR(&struct_unpack_obj) },
     { MP_ROM_QSTR(MP_QSTR_unpack_from), MP_ROM_PTR(&struct_unpack_from_obj) },
 };