From 3d9a39e211ed98033b6d6fcda5d0771bfb25dc46 Mon Sep 17 00:00:00 2001
From: xbe <xbe@machine>
Date: Tue, 8 Apr 2014 11:42:19 -0700
Subject: [PATCH] py: Implement str.[r]index() and add tests for them.

---
 py/objstr.c                   | 24 +++++++++--
 py/qstrdefs.h                 |  1 +
 tests/basics/string_index.py  | 78 +++++++++++++++++++++++++++++++++++
 tests/basics/string_rindex.py | 78 +++++++++++++++++++++++++++++++++++
 4 files changed, 177 insertions(+), 4 deletions(-)
 create mode 100644 tests/basics/string_index.py
 create mode 100644 tests/basics/string_rindex.py

diff --git a/py/objstr.c b/py/objstr.c
index 7af08fe0b..e37d7b22f 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -441,7 +441,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
     return res;
 }
 
-STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction) {
+STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction, bool is_index) {
     assert(2 <= n_args && n_args <= 4);
     assert(MP_OBJ_IS_STR(args[0]));
     assert(MP_OBJ_IS_STR(args[1]));
@@ -461,7 +461,11 @@ STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t dire
     const byte *p = find_subbytes(haystack + start, end - start, needle, needle_len, direction);
     if (p == NULL) {
         // not found
-        return MP_OBJ_NEW_SMALL_INT(-1);
+        if (is_index) {
+            nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "substring not found"));
+        } else {
+            return MP_OBJ_NEW_SMALL_INT(-1);
+        }
     } else {
         // found
         return MP_OBJ_NEW_SMALL_INT(p - haystack);
@@ -469,11 +473,19 @@ STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t dire
 }
 
 STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
-    return str_finder(n_args, args, 1);
+    return str_finder(n_args, args, 1, false);
 }
 
 STATIC mp_obj_t str_rfind(uint n_args, const mp_obj_t *args) {
-    return str_finder(n_args, args, -1);
+    return str_finder(n_args, args, -1, false);
+}
+
+STATIC mp_obj_t str_index(uint n_args, const mp_obj_t *args) {
+    return str_finder(n_args, args, 1, true);
+}
+
+STATIC mp_obj_t str_rindex(uint n_args, const mp_obj_t *args) {
+    return str_finder(n_args, args, -1, true);
 }
 
 // TODO: (Much) more variety in args
@@ -1307,6 +1319,8 @@ STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, in
 
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rfind_obj, 2, 4, str_rfind);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_index_obj, 2, 4, str_index);
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rindex_obj, 2, 4, str_rindex);
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
@@ -1320,6 +1334,8 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition);
 STATIC const mp_map_elem_t str_locals_dict_table[] = {
     { MP_OBJ_NEW_QSTR(MP_QSTR_find), (mp_obj_t)&str_find_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_rfind), (mp_obj_t)&str_rfind_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR_index), (mp_obj_t)&str_index_obj },
+    { MP_OBJ_NEW_QSTR(MP_QSTR_rindex), (mp_obj_t)&str_rindex_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_join), (mp_obj_t)&str_join_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_split), (mp_obj_t)&str_split_obj },
     { MP_OBJ_NEW_QSTR(MP_QSTR_startswith), (mp_obj_t)&str_startswith_obj },
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 29eb51d0b..a379da391 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -181,6 +181,7 @@ Q(union)
 Q(update)
 Q(find)
 Q(rfind)
+Q(rindex)
 Q(split)
 Q(startswith)
 Q(replace)
diff --git a/tests/basics/string_index.py b/tests/basics/string_index.py
new file mode 100644
index 000000000..31f6900e6
--- /dev/null
+++ b/tests/basics/string_index.py
@@ -0,0 +1,78 @@
+print("hello world".index("ll"))
+print("hello world".index("ll", None))
+print("hello world".index("ll", 1))
+print("hello world".index("ll", 1, None))
+print("hello world".index("ll", None, None))
+print("hello world".index("ll", 1, -1))
+
+try:
+    print("hello world".index("ll", 1, 1))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("hello world".index("ll", 1, 2))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("hello world".index("ll", 1, 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+print("hello world".index("ll", 1, 4))
+print("hello world".index("ll", 1, 5))
+print("hello world".index("ll", -100))
+print("0000".index('0'))
+print("0000".index('0', 0))
+print("0000".index('0', 1))
+print("0000".index('0', 2))
+print("0000".index('0', 3))
+
+try:
+    print("0000".index('0', 4))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".index('0', 5))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".index('-1', 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".index('1', 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".index('1', 4))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".index('1', 5))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
diff --git a/tests/basics/string_rindex.py b/tests/basics/string_rindex.py
new file mode 100644
index 000000000..25acd37f6
--- /dev/null
+++ b/tests/basics/string_rindex.py
@@ -0,0 +1,78 @@
+print("hello world".rindex("ll"))
+print("hello world".rindex("ll", None))
+print("hello world".rindex("ll", 1))
+print("hello world".rindex("ll", 1, None))
+print("hello world".rindex("ll", None, None))
+print("hello world".rindex("ll", 1, -1))
+
+try:
+    print("hello world".rindex("ll", 1, 1))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("hello world".rindex("ll", 1, 2))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("hello world".rindex("ll", 1, 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+print("hello world".rindex("ll", 1, 4))
+print("hello world".rindex("ll", 1, 5))
+print("hello world".rindex("ll", -100))
+print("0000".rindex('0'))
+print("0000".rindex('0', 0))
+print("0000".rindex('0', 1))
+print("0000".rindex('0', 2))
+print("0000".rindex('0', 3))
+
+try:
+    print("0000".rindex('0', 4))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".rindex('0', 5))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".rindex('-1', 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".rindex('1', 3))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".rindex('1', 4))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
+
+try:
+    print("0000".rindex('1', 5))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")
-- 
GitLab