From 613a8e3edf078c284bd981426cc5a256eabb2323 Mon Sep 17 00:00:00 2001
From: xbe <xbe@machine>
Date: Tue, 18 Mar 2014 00:06:29 -0700
Subject: [PATCH] Implement str.partition and add tests for it.

---
 py/objstr.c                      | 27 +++++++++++++++++++++++++++
 tests/basics/string_partition.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)
 create mode 100644 tests/basics/string_partition.py

diff --git a/py/objstr.c b/py/objstr.c
index d660bf952b..03711debb2 100644
--- a/py/objstr.c
+++ b/py/objstr.c
@@ -520,6 +520,31 @@ STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
     return MP_OBJ_NEW_SMALL_INT(num_occurrences);
 }
 
+STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
+    assert(MP_OBJ_IS_STR(self_in));
+    if (!MP_OBJ_IS_STR(arg)) {
+        nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
+                                               "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
+    }
+
+    GET_STR_DATA_LEN(self_in, str, str_len);
+    GET_STR_DATA_LEN(arg, sep, sep_len);
+
+    if (sep_len == 0) {
+        nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
+    }
+
+    for (machine_uint_t str_index = 0; str_index + sep_len <= str_len; str_index++) {
+        if (memcmp(&str[str_index], sep, sep_len) == 0) {
+            mp_obj_t items[] = {mp_obj_new_str(str, str_index, false), arg,
+                                mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false)};
+            return mp_obj_new_tuple(3, items);
+        }
+    }
+    mp_obj_t items[] = {mp_obj_new_str(str, str_len, false), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
+    return mp_obj_new_tuple(3, items);
+}
+
 STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
     if (flags == BUFFER_READ) {
         GET_STR_DATA_LEN(self_in, str_data, str_len);
@@ -542,6 +567,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
 STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
+STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_partition_obj, str_partition);
 
 STATIC const mp_method_t str_type_methods[] = {
     { "find", &str_find_obj },
@@ -552,6 +578,7 @@ STATIC const mp_method_t str_type_methods[] = {
     { "format", &str_format_obj },
     { "replace", &str_replace_obj },
     { "count", &str_count_obj },
+    { "partition", &str_partition_obj },
     { NULL, NULL }, // end-of-list sentinel
 };
 
diff --git a/tests/basics/string_partition.py b/tests/basics/string_partition.py
new file mode 100644
index 0000000000..ad70d02509
--- /dev/null
+++ b/tests/basics/string_partition.py
@@ -0,0 +1,29 @@
+print("asdf".partition('g'))
+print("asdf".partition('a'))
+print("asdf".partition('s'))
+print("asdf".partition('f'))
+print("asdf".partition('d'))
+print("asdf".partition('asd'))
+print("asdf".partition('sdf'))
+print("asdf".partition('as'))
+print("asdf".partition('df'))
+print("asdf".partition('asdf'))
+print("asdf".partition('asdfa'))
+print("asdf".partition('fasdf'))
+print("asdf".partition('fasdfa'))
+print("abba".partition('a'))
+print("abba".partition('b'))
+
+try:
+    print("asdf".partition(1))
+except TypeError:
+    print("Raised TypeError")
+else:
+    print("Did not raise TypeError")
+
+try:
+    print("asdf".partition(''))
+except ValueError:
+    print("Raised ValueError")
+else:
+    print("Did not raise ValueError")