objstr: Make .split() support bytes.

pfalcon · pfalcon · commit bfb8819c0c7e · 2014-05-11T21:17:28.000+03:00
diff --git a/py/objstr.c b/py/objstr.c
@@ -411,6 +411,7 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
 #define is_ws(c) ((c) == ' ' || (c) == '\t')
 
 STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
+    const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);
     machine_int_t splits = -1;
     mp_obj_t sep = mp_const_none;
     if (n_args > 1) {
@@ -432,7 +433,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
         while (s < top && splits != 0) {
             const byte *start = s;
             while (s < top && !is_ws(*s)) s++;
-            mp_obj_list_append(res, mp_obj_new_str(start, s - start, false));
+            mp_obj_list_append(res, str_new(self_type, start, s - start));
             if (s >= top) {
                 break;
             }
@@ -443,7 +444,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
         }
 
         if (s < top) {
-            mp_obj_list_append(res, mp_obj_new_str(s, top - s, false));
+            mp_obj_list_append(res, str_new(self_type, s, top - s));
         }
 
     } else {
@@ -467,7 +468,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
                 }
                 s++;
             }
-            mp_obj_list_append(res, mp_obj_new_str(start, s - start, false));
+            mp_obj_list_append(res, str_new(self_type, start, s - start));
             if (s >= top) {
                 break;
             }
diff --git a/tests/basics/string_split.py b/tests/basics/string_split.py
@@ -26,3 +26,5 @@
 print("abcabc".split("bc", 0))
 print("abcabc".split("bc", 1))
 print("abcabc".split("bc", 2))
+
+print(b"abcabc".split(b"bc", 2))

Original file line number	Diff line number	Diff line change
`@@ -411,6 +411,7 @@ STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {`
`411`	`411`	`#define is_ws(c) ((c) == ' ' \|\| (c) == '\t')`
`412`	`412`
`413`	`413`	`STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {`
	`414`	`+ const mp_obj_type_t *self_type = mp_obj_get_type(args[0]);`
`414`	`415`	`machine_int_t splits = -1;`
`415`	`416`	`mp_obj_t sep = mp_const_none;`
`416`	`417`	`if (n_args > 1) {`
`@@ -432,7 +433,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {`
`432`	`433`	`while (s < top && splits != 0) {`
`433`	`434`	`const byte *start = s;`
`434`	`435`	`while (s < top && !is_ws(*s)) s++;`
`435`		`- mp_obj_list_append(res, mp_obj_new_str(start, s - start, false));`
	`436`	`+ mp_obj_list_append(res, str_new(self_type, start, s - start));`
`436`	`437`	`if (s >= top) {`
`437`	`438`	`break;`
`438`	`439`	`}`
`@@ -443,7 +444,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {`
`443`	`444`	`}`
`444`	`445`
`445`	`446`	`if (s < top) {`
`446`		`- mp_obj_list_append(res, mp_obj_new_str(s, top - s, false));`
	`447`	`+ mp_obj_list_append(res, str_new(self_type, s, top - s));`
`447`	`448`	`}`
`448`	`449`
`449`	`450`	`} else {`
`@@ -467,7 +468,7 @@ STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {`
`467`	`468`	`}`
`468`	`469`	`s++;`
`469`	`470`	`}`
`470`		`- mp_obj_list_append(res, mp_obj_new_str(start, s - start, false));`
	`471`	`+ mp_obj_list_append(res, str_new(self_type, start, s - start));`
`471`	`472`	`if (s >= top) {`
`472`	`473`	`break;`
`473`	`474`	`}`