Skip to content

Commit c38809e

Browse files
committed
py/objarray: Implement "in" operator for bytearray.
1 parent 609a9c6 commit c38809e

3 files changed

Lines changed: 26 additions & 1 deletion

File tree

py/objarray.c

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
#include "py/runtime0.h"
3434
#include "py/runtime.h"
3535
#include "py/binary.h"
36+
#include "py/objstr.h"
3637

3738
#if MICROPY_PY_ARRAY || MICROPY_PY_BUILTINS_BYTEARRAY || MICROPY_PY_BUILTINS_MEMORYVIEW
3839

@@ -283,6 +284,29 @@ STATIC mp_obj_t array_binary_op(mp_uint_t op, mp_obj_t lhs_in, mp_obj_t rhs_in)
283284
return lhs_in;
284285
}
285286

287+
case MP_BINARY_OP_IN: {
288+
/* NOTE `a in b` is `b.__contains__(a)` */
289+
mp_buffer_info_t lhs_bufinfo;
290+
mp_buffer_info_t rhs_bufinfo;
291+
292+
// Can search string only in bytearray
293+
if (mp_get_buffer(rhs_in, &rhs_bufinfo, MP_BUFFER_READ)) {
294+
if (!MP_OBJ_IS_TYPE(lhs_in, &mp_type_bytearray)) {
295+
return mp_const_false;
296+
}
297+
array_get_buffer(lhs_in, &lhs_bufinfo, MP_BUFFER_READ);
298+
return mp_obj_new_bool(
299+
find_subbytes(lhs_bufinfo.buf, lhs_bufinfo.len, rhs_bufinfo.buf, rhs_bufinfo.len, 1) != NULL);
300+
}
301+
302+
// Otherwise, can only look for a scalar numeric value in an array
303+
if (MP_OBJ_IS_INT(rhs_in) || mp_obj_is_float(rhs_in)) {
304+
mp_not_implemented("");
305+
}
306+
307+
return mp_const_false;
308+
}
309+
286310
case MP_BINARY_OP_EQUAL: {
287311
mp_buffer_info_t lhs_bufinfo;
288312
mp_buffer_info_t rhs_bufinfo;

py/objstr.c

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,7 @@ STATIC mp_obj_t bytes_make_new(const mp_obj_type_t *type_in, size_t n_args, size
245245

246246
// like strstr but with specified length and allows \0 bytes
247247
// TODO replace with something more efficient/standard
248-
STATIC const byte *find_subbytes(const byte *haystack, mp_uint_t hlen, const byte *needle, mp_uint_t nlen, mp_int_t direction) {
248+
const byte *find_subbytes(const byte *haystack, mp_uint_t hlen, const byte *needle, mp_uint_t nlen, mp_int_t direction) {
249249
if (hlen >= nlen) {
250250
mp_uint_t str_index, str_index_end;
251251
if (direction > 0) {

py/objstr.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ mp_int_t mp_obj_str_get_buffer(mp_obj_t self_in, mp_buffer_info_t *bufinfo, mp_u
7171

7272
const byte *str_index_to_ptr(const mp_obj_type_t *type, const byte *self_data, size_t self_len,
7373
mp_obj_t index, bool is_slice);
74+
const byte *find_subbytes(const byte *haystack, mp_uint_t hlen, const byte *needle, mp_uint_t nlen, mp_int_t direction);
7475

7576
MP_DECLARE_CONST_FUN_OBJ(str_encode_obj);
7677
MP_DECLARE_CONST_FUN_OBJ(str_find_obj);

0 commit comments

Comments
 (0)