Merge pull request #3161 from tannewt/bytearray_find

Add find variants to bytearray

3 files changed, 78 insertions, 5 deletions

diff --git a/locale/circuitpython.pot b/locale/circuitpython.pot
index 6bdc91bf1..988f89a05 100644
--- a/locale/circuitpython.pot
+++ b/locale/circuitpython.pot
@@ -1,12 +1,14 @@
-# SPDX-FileCopyrightText: 2014 MicroPython & CircuitPython contributors (https://github.com/adafruit/circuitpython/graphs/contributors)
+# SOME DESCRIPTIVE TITLE.
+# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER
+# This file is distributed under the same license as the PACKAGE package.
+# FIRST AUTHOR <EMAIL@ADDRESS>, YEAR.
 #
-# SPDX-License-Identifier: MIT
-
+#, fuzzy
 msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2020-07-07 14:38-0500\n"
+"POT-Creation-Date: 2020-07-17 18:03-0700\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <LL@li.org>\n"
@@ -103,6 +105,10 @@ msgstr ""
 msgid "'%q' argument required"
 msgstr ""
 
+#: py/objarray.c
+msgid "'%q' object is not bytes-like"
+msgstr ""
+
 #: py/emitinlinethumb.c py/emitinlinextensa.c
 #, c-format
 msgid "'%s' expects a label"
@@ -3098,7 +3104,7 @@ msgstr ""
 msgid "struct: no fields"
 msgstr ""
 
-#: py/objstr.c
+#: py/objarray.c py/objstr.c
 msgid "substring not found"
 msgstr ""
 
diff --git a/ports/atmel-samd/boards/circuitbrains_basic_m0/mpconfigboard.mk b/ports/atmel-samd/boards/circuitbrains_basic_m0/mpconfigboard.mk
index da5f1ffac..619b7d694 100755
--- a/ports/atmel-samd/boards/circuitbrains_basic_m0/mpconfigboard.mk
+++ b/ports/atmel-samd/boards/circuitbrains_basic_m0/mpconfigboard.mk
@@ -12,6 +12,7 @@ EXTERNAL_FLASH_DEVICES = "W25Q32JV_IQ"
 LONGINT_IMPL = MPZ
 
 CIRCUITPY_BITBANGIO = 0
+CIRCUITPY_COUNTIO = 0
 CIRCUITPY_FREQUENCYIO = 0
 CIRCUITPY_I2CPERIPHERAL = 0
 CIRCUITPY_VECTORIO = 0
diff --git a/py/objarray.c b/py/objarray.c
index 7dfdc5b12..5d83f0697 100644
--- a/py/objarray.c
+++ b/py/objarray.c
@@ -400,6 +400,66 @@ 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 && MICROPY_CPYTHON_COMPAT
+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) {
+        mp_raise_TypeError(translate("a bytes-like object is required"));
+    }
+
+    const byte *start = haystack_bufinfo.buf;
+    const byte *end = ((const byte*)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 = ((const byte*)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,7 +640,13 @@ 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) },
+
 #if MICROPY_CPYTHON_COMPAT
+    { 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) },
+
     { MP_ROM_QSTR(MP_QSTR_decode), MP_ROM_PTR(&array_decode_obj) },
 #endif
 };