Skip to content

Commit a17755e

Browse files
dhylandsdpgeorge
dhylands
authored and
dpgeorge
committed
py: Add ustruct.pack_into and unpack_from
1 parent ac16cc9 commit a17755e

3 files changed

Lines changed: 118 additions & 19 deletions

File tree

‎py/modstruct.c‎

Lines changed: 75 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -132,22 +132,45 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
132132
}
133133
MP_DEFINE_CONST_FUN_OBJ_1(struct_calcsize_obj, struct_calcsize);
134134

135-
STATIC mp_obj_t struct_unpack(mp_obj_t fmt_in, mp_obj_t data_in) {
136-
// TODO: "The buffer must contain exactly the amount of data required by the format (len(bytes) must equal calcsize(fmt))."
137-
const char *fmt = mp_obj_str_get_str(fmt_in);
135+
STATIC mp_obj_t struct_unpack_from(size_t n_args, const mp_obj_t *args) {
136+
// unpack requires that the buffer be exactly the right size.
137+
// unpack_from requires that the buffer be "big enough".
138+
// Since we implement unpack and unpack_from using the same function
139+
// we relax the "exact" requirement, and only implement "big enough".
140+
const char *fmt = mp_obj_str_get_str(args[0]);
138141
char fmt_type = get_fmt_type(&fmt);
139-
uint size = calcsize_items(fmt);
140-
mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(size, NULL));
142+
uint num_items = calcsize_items(fmt);
143+
mp_obj_tuple_t *res = MP_OBJ_TO_PTR(mp_obj_new_tuple(num_items, NULL));
141144
mp_buffer_info_t bufinfo;
142-
mp_get_buffer_raise(data_in, &bufinfo, MP_BUFFER_READ);
145+
mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_READ);
143146
byte *p = bufinfo.buf;
147+
byte *end_p = &p[bufinfo.len];
148+
mp_int_t offset = 0;
144149

145-
for (uint i = 0; i < size;) {
150+
if (n_args > 2) {
151+
// offset arg provided
152+
offset = mp_obj_get_int(args[2]);
153+
if (offset < 0) {
154+
// negative offsets are relative to the end of the buffer
155+
offset = bufinfo.len + offset;
156+
if (offset < 0) {
157+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
158+
}
159+
}
160+
p += offset;
161+
}
162+
163+
for (uint i = 0; i < num_items;) {
164+
if (*fmt == '\0') {
165+
break;
166+
}
146167
mp_uint_t sz = 1;
147168
if (unichar_isdigit(*fmt)) {
148169
sz = get_fmt_num(&fmt);
149170
}
150-
171+
if (p + sz > end_p) {
172+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
173+
}
151174
mp_obj_t item;
152175
if (*fmt == 's') {
153176
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) {
163186
}
164187
return MP_OBJ_FROM_PTR(res);
165188
}
166-
MP_DEFINE_CONST_FUN_OBJ_2(struct_unpack_obj, struct_unpack);
189+
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_unpack_from_obj, 2, 3, struct_unpack_from);
167190

168-
STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
169-
// TODO: "The arguments must match the values required by the format exactly."
170-
const char *fmt = mp_obj_str_get_str(args[0]);
191+
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) {
192+
const char *fmt = mp_obj_str_get_str(fmt_in);
171193
char fmt_type = get_fmt_type(&fmt);
172-
mp_int_t size = MP_OBJ_SMALL_INT_VALUE(struct_calcsize(args[0]));
173-
vstr_t vstr;
174-
vstr_init_len(&vstr, size);
175-
byte *p = (byte*)vstr.buf;
176-
memset(p, 0, size);
177194

178-
for (mp_uint_t i = 1; i < n_args;) {
195+
size_t i;
196+
for (i = 0; i < n_args;) {
179197
mp_uint_t sz = 1;
198+
if (*fmt == '\0') {
199+
break;
200+
}
180201
if (unichar_isdigit(*fmt)) {
181202
sz = get_fmt_num(&fmt);
182203
}
204+
if (p + sz > end_p) {
205+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
206+
}
183207

184208
if (*fmt == 's') {
185209
mp_buffer_info_t bufinfo;
@@ -198,16 +222,48 @@ STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
198222
}
199223
fmt++;
200224
}
225+
}
201226

227+
STATIC mp_obj_t struct_pack(size_t n_args, const mp_obj_t *args) {
228+
// TODO: "The arguments must match the values required by the format exactly."
229+
mp_int_t size = MP_OBJ_SMALL_INT_VALUE(struct_calcsize(args[0]));
230+
vstr_t vstr;
231+
vstr_init_len(&vstr, size);
232+
byte *p = (byte*)vstr.buf;
233+
memset(p, 0, size);
234+
byte *end_p = &p[size];
235+
struct_pack_into_internal(args[0], p, end_p, n_args - 1, &args[1]);
202236
return mp_obj_new_str_from_vstr(&mp_type_bytes, &vstr);
203237
}
204238
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_obj, 1, MP_OBJ_FUN_ARGS_MAX, struct_pack);
205239

240+
STATIC mp_obj_t struct_pack_into(size_t n_args, const mp_obj_t *args) {
241+
mp_buffer_info_t bufinfo;
242+
mp_get_buffer_raise(args[1], &bufinfo, MP_BUFFER_WRITE);
243+
mp_int_t offset = mp_obj_get_int(args[2]);
244+
if (offset < 0) {
245+
// negative offsets are relative to the end of the buffer
246+
offset = (mp_int_t)bufinfo.len + offset;
247+
if (offset < 0) {
248+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "buffer too small"));
249+
}
250+
}
251+
byte *p = (byte *)bufinfo.buf;
252+
byte *end_p = &p[bufinfo.len];
253+
p += offset;
254+
255+
struct_pack_into_internal(args[0], p, end_p, n_args - 3, &args[3]);
256+
return mp_const_none;
257+
}
258+
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(struct_pack_into_obj, 3, MP_OBJ_FUN_ARGS_MAX, struct_pack_into);
259+
206260
STATIC const mp_rom_map_elem_t mp_module_struct_globals_table[] = {
207261
{ MP_ROM_QSTR(MP_QSTR___name__), MP_ROM_QSTR(MP_QSTR_ustruct) },
208262
{ MP_ROM_QSTR(MP_QSTR_calcsize), MP_ROM_PTR(&struct_calcsize_obj) },
209263
{ MP_ROM_QSTR(MP_QSTR_pack), MP_ROM_PTR(&struct_pack_obj) },
210-
{ MP_ROM_QSTR(MP_QSTR_unpack), MP_ROM_PTR(&struct_unpack_obj) },
264+
{ MP_ROM_QSTR(MP_QSTR_pack_into), MP_ROM_PTR(&struct_pack_into_obj) },
265+
{ MP_ROM_QSTR(MP_QSTR_unpack), MP_ROM_PTR(&struct_unpack_from_obj) },
266+
{ MP_ROM_QSTR(MP_QSTR_unpack_from), MP_ROM_PTR(&struct_unpack_from_obj) },
211267
};
212268

213269
STATIC MP_DEFINE_CONST_DICT(mp_module_struct_globals, mp_module_struct_globals_table);

‎py/qstrdefs.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,7 +491,9 @@ Q(print_exception)
491491
Q(struct)
492492
Q(ustruct)
493493
Q(pack)
494+
Q(pack_into)
494495
Q(unpack)
496+
Q(unpack_from)
495497
Q(calcsize)
496498
#endif
497499

‎tests/basics/struct1.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,3 +66,44 @@
6666
# but later were implemented for all.
6767
print(struct.unpack("<3B2h", b"foo\x12\x34\xff\xff"))
6868
print(struct.pack("<3B", 1, 2, 3))
69+
70+
# pack_into
71+
buf = bytearray(b'>>>123<<<')
72+
struct.pack_into('<bbb', buf, 3, 0x41, 0x42, 0x43)
73+
print(buf)
74+
struct.pack_into('<bbb', buf, -6, 0x44, 0x45, 0x46)
75+
print(buf)
76+
77+
try:
78+
struct.pack_into('<bbb', buf, 7, 0x41, 0x42, 0x43)
79+
except:
80+
print('struct.error')
81+
try:
82+
struct.pack_into('<bbb', buf, -10, 0x41, 0x42, 0x43)
83+
except:
84+
print('struct.error')
85+
86+
# unpack_from
87+
buf = b'0123456789'
88+
print(struct.unpack_from('<b', buf, 4))
89+
print(struct.unpack_from('<b', buf, -4))
90+
try:
91+
print(struct.unpack_from('<b', buf, 10))
92+
except:
93+
print('struct.error')
94+
try:
95+
print(struct.unpack_from('<b', buf, -11))
96+
except:
97+
print('struct.error')
98+
99+
# pack with too many args, not checked by uPy
100+
#try:
101+
# print(struct.pack('ii', 1, 2, 3))
102+
#except:
103+
# print('struct.error')
104+
105+
# pack with too few args, not checked by uPy
106+
#try:
107+
# print(struct.pack('ii', 1))
108+
#except:
109+
# print('struct.error')

0 commit comments

Comments
 (0)