Add find varients to bytearray

This commit is contained in:
Scott Shawcroft 2020-07-15 17:59:11 -07:00
parent 9cdf5e148a
commit 4b157aa6b8
No known key found for this signature in database
GPG Key ID: 9349BC7E64B1921E
1 changed files with 70 additions and 0 deletions

View File

@ -400,6 +400,70 @@ STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in) {
STATIC MP_DEFINE_CONST_FUN_OBJ_2(array_extend_obj, array_extend);
#endif
#if MICROPY_PY_BUILTINS_BYTEARRAY
STATIC mp_obj_t buffer_finder(size_t n_args, const mp_obj_t *args, int direction, bool is_index) {
mp_check_self(MP_OBJ_IS_TYPE(args[0], &mp_type_bytearray));
const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);
mp_buffer_info_t haystack_bufinfo;
mp_get_buffer_raise(args[0], &haystack_bufinfo, MP_BUFFER_READ);
mp_buffer_info_t needle_bufinfo;
mp_get_buffer_raise(args[1], &needle_bufinfo, MP_BUFFER_READ);
if (mp_binary_get_size('@', needle_bufinfo.typecode, NULL) != 1) {
const qstr src_name = mp_obj_get_type(args[1])->name;
nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
translate("'%q' object is not bytes-like"),
src_name));
}
const byte *start = haystack_bufinfo.buf;
const byte *end = haystack_bufinfo.buf + haystack_bufinfo.len;
if (n_args >= 3 && args[2] != mp_const_none) {
start += mp_get_index(self_type, haystack_bufinfo.len, args[2], true);
}
if (n_args >= 4 && args[3] != mp_const_none) {
end = haystack_bufinfo.buf + mp_get_index(self_type, haystack_bufinfo.len, args[3], true);
}
const byte *p = NULL;
if (end >= start) {
p = find_subbytes(start, end - start, needle_bufinfo.buf, needle_bufinfo.len, direction);
}
if (p == NULL) {
if (is_index) {
mp_raise_ValueError(translate("substring not found"));
} else {
return MP_OBJ_NEW_SMALL_INT(-1);
}
}
return MP_OBJ_NEW_SMALL_INT(p - (const byte*) haystack_bufinfo.buf);
}
STATIC mp_obj_t buffer_find(size_t n_args, const mp_obj_t *args) {
return buffer_finder(n_args, args, 1, false);
}
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(buffer_find_obj, 2, 4, buffer_find);
STATIC mp_obj_t buffer_rfind(size_t n_args, const mp_obj_t *args) {
return buffer_finder(n_args, args, -1, false);
}
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(buffer_rfind_obj, 2, 4, buffer_rfind);
STATIC mp_obj_t buffer_index(size_t n_args, const mp_obj_t *args) {
return buffer_finder(n_args, args, 1, true);
}
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(buffer_index_obj, 2, 4, buffer_index);
STATIC mp_obj_t buffer_rindex(size_t n_args, const mp_obj_t *args) {
return buffer_finder(n_args, args, -1, true);
}
MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(buffer_rindex_obj, 2, 4, buffer_rindex);
#endif
STATIC mp_obj_t array_subscr(mp_obj_t self_in, mp_obj_t index_in, mp_obj_t value) {
if (value == MP_OBJ_NULL) {
// delete item
@ -580,6 +644,12 @@ STATIC MP_DEFINE_CONST_DICT(array_locals_dict, array_locals_dict_table);
STATIC const mp_rom_map_elem_t bytearray_locals_dict_table[] = {
{ MP_ROM_QSTR(MP_QSTR_append), MP_ROM_PTR(&array_append_obj) },
{ MP_ROM_QSTR(MP_QSTR_extend), MP_ROM_PTR(&array_extend_obj) },
{ MP_ROM_QSTR(MP_QSTR_find), MP_ROM_PTR(&buffer_find_obj) },
{ MP_ROM_QSTR(MP_QSTR_rfind), MP_ROM_PTR(&buffer_rfind_obj) },
{ MP_ROM_QSTR(MP_QSTR_index), MP_ROM_PTR(&buffer_index_obj) },
{ MP_ROM_QSTR(MP_QSTR_rindex), MP_ROM_PTR(&buffer_rindex_obj) },
#if MICROPY_CPYTHON_COMPAT
{ MP_ROM_QSTR(MP_QSTR_decode), MP_ROM_PTR(&array_decode_obj) },
#endif