blob: c719935789b24d346bbe3d6ac35322f69a48a91e [file] [log] [blame]
xbeefe34222014-03-16 00:14:26 -07001#include <stdbool.h>
Damiend99b0522013-12-21 18:17:45 +00002#include <string.h>
3#include <assert.h>
4
5#include "nlr.h"
6#include "misc.h"
7#include "mpconfig.h"
Damien George55baff42014-01-21 21:40:13 +00008#include "qstr.h"
Damiend99b0522013-12-21 18:17:45 +00009#include "obj.h"
10#include "runtime0.h"
11#include "runtime.h"
12
13typedef struct _mp_obj_str_t {
14 mp_obj_base_t base;
Damien George5fa93b62014-01-22 14:35:10 +000015 machine_uint_t hash : 16; // XXX here we assume the hash size is 16 bits (it is at the moment; see qstr.c)
16 machine_uint_t len : 16; // len == number of bytes used in data, alloc = len + 1 because (at the moment) we also append a null byte
17 byte data[];
Damiend99b0522013-12-21 18:17:45 +000018} mp_obj_str_t;
19
Damien George5fa93b62014-01-22 14:35:10 +000020// use this macro to extract the string hash
21#define GET_STR_HASH(str_obj_in, str_hash) uint str_hash; if (MP_OBJ_IS_QSTR(str_obj_in)) { str_hash = qstr_hash(MP_OBJ_QSTR_VALUE(str_obj_in)); } else { str_hash = ((mp_obj_str_t*)str_obj_in)->hash; }
22
23// use this macro to extract the string length
24#define GET_STR_LEN(str_obj_in, str_len) uint str_len; if (MP_OBJ_IS_QSTR(str_obj_in)) { str_len = qstr_len(MP_OBJ_QSTR_VALUE(str_obj_in)); } else { str_len = ((mp_obj_str_t*)str_obj_in)->len; }
25
26// use this macro to extract the string data and length
27#define GET_STR_DATA_LEN(str_obj_in, str_data, str_len) const byte *str_data; uint str_len; if (MP_OBJ_IS_QSTR(str_obj_in)) { str_data = qstr_data(MP_OBJ_QSTR_VALUE(str_obj_in), &str_len); } else { str_len = ((mp_obj_str_t*)str_obj_in)->len; str_data = ((mp_obj_str_t*)str_obj_in)->data; }
28
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020029STATIC mp_obj_t mp_obj_new_str_iterator(mp_obj_t str);
30STATIC mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str);
xyb8cfc9f02014-01-05 18:47:51 +080031
32/******************************************************************************/
33/* str */
34
Paul Sokolovsky0b7e29c2014-01-28 03:40:06 +020035void mp_str_print_quoted(void (*print)(void *env, const char *fmt, ...), void *env, const byte *str_data, uint str_len) {
36 // this escapes characters, but it will be very slow to print (calling print many times)
37 bool has_single_quote = false;
38 bool has_double_quote = false;
39 for (const byte *s = str_data, *top = str_data + str_len; (!has_single_quote || !has_double_quote) && s < top; s++) {
40 if (*s == '\'') {
41 has_single_quote = true;
42 } else if (*s == '"') {
43 has_double_quote = true;
44 }
45 }
46 int quote_char = '\'';
47 if (has_single_quote && !has_double_quote) {
48 quote_char = '"';
49 }
50 print(env, "%c", quote_char);
51 for (const byte *s = str_data, *top = str_data + str_len; s < top; s++) {
52 if (*s == quote_char) {
53 print(env, "\\%c", quote_char);
54 } else if (*s == '\\') {
55 print(env, "\\\\");
56 } else if (32 <= *s && *s <= 126) {
57 print(env, "%c", *s);
58 } else if (*s == '\n') {
59 print(env, "\\n");
60 // TODO add more escape codes here if we want to match CPython
61 } else {
62 print(env, "\\x%02x", *s);
63 }
64 }
65 print(env, "%c", quote_char);
66}
67
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020068STATIC void str_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
Damien George5fa93b62014-01-22 14:35:10 +000069 GET_STR_DATA_LEN(self_in, str_data, str_len);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +020070 bool is_bytes = MP_OBJ_IS_TYPE(self_in, &bytes_type);
71 if (kind == PRINT_STR && !is_bytes) {
Damien George5fa93b62014-01-22 14:35:10 +000072 print(env, "%.*s", str_len, str_data);
Paul Sokolovsky76d982e2014-01-13 19:19:16 +020073 } else {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +020074 if (is_bytes) {
75 print(env, "b");
76 }
Paul Sokolovsky0b7e29c2014-01-28 03:40:06 +020077 mp_str_print_quoted(print, env, str_data, str_len);
Paul Sokolovsky76d982e2014-01-13 19:19:16 +020078 }
Damiend99b0522013-12-21 18:17:45 +000079}
80
Damien George55baff42014-01-21 21:40:13 +000081// like strstr but with specified length and allows \0 bytes
82// TODO replace with something more efficient/standard
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020083STATIC const byte *find_subbytes(const byte *haystack, uint hlen, const byte *needle, uint nlen) {
Damien George55baff42014-01-21 21:40:13 +000084 if (hlen >= nlen) {
85 for (uint i = 0; i <= hlen - nlen; i++) {
86 bool found = true;
87 for (uint j = 0; j < nlen; j++) {
88 if (haystack[i + j] != needle[j]) {
89 found = false;
90 break;
91 }
92 }
93 if (found) {
94 return haystack + i;
95 }
96 }
97 }
98 return NULL;
99}
100
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200101STATIC mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
Damien George5fa93b62014-01-22 14:35:10 +0000102 GET_STR_DATA_LEN(lhs_in, lhs_data, lhs_len);
Damiend99b0522013-12-21 18:17:45 +0000103 switch (op) {
104 case RT_BINARY_OP_SUBSCR:
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200105 // TODO: need predicate to check for int-like type (bools are such for example)
106 // ["no", "yes"][1 == 2] is common idiom
107 if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
xbe9e1e8cd2014-03-12 22:57:16 -0700108 uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200109 if (MP_OBJ_IS_TYPE(lhs_in, &bytes_type)) {
Damien George7c9c6672014-01-25 00:17:36 +0000110 return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200111 } else {
112 return mp_obj_new_str(lhs_data + index, 1, true);
113 }
Paul Sokolovskye606cb62014-01-04 01:34:23 +0200114#if MICROPY_ENABLE_SLICE
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200115 } else if (MP_OBJ_IS_TYPE(rhs_in, &slice_type)) {
Paul Sokolovsky7364af22014-02-02 02:38:22 +0200116 machine_uint_t start, stop;
Paul Sokolovskyea2509d2014-02-02 08:57:05 +0200117 if (!m_seq_get_fast_slice_indexes(lhs_len, rhs_in, &start, &stop)) {
118 assert(0);
119 }
Damien George5fa93b62014-01-22 14:35:10 +0000120 return mp_obj_new_str(lhs_data + start, stop - start, false);
Paul Sokolovskye606cb62014-01-04 01:34:23 +0200121#endif
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200122 } else {
Paul Sokolovskyf8b9d3c2014-01-04 01:38:26 +0200123 // Message doesn't match CPython, but we don't have so much bytes as they
124 // to spend them on verbose wording
Damien Georgec5966122014-02-15 16:10:44 +0000125 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "index must be int"));
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200126 }
Damiend99b0522013-12-21 18:17:45 +0000127
128 case RT_BINARY_OP_ADD:
129 case RT_BINARY_OP_INPLACE_ADD:
Damien George5fa93b62014-01-22 14:35:10 +0000130 if (MP_OBJ_IS_STR(rhs_in)) {
Damiend99b0522013-12-21 18:17:45 +0000131 // add 2 strings
Damien George5fa93b62014-01-22 14:35:10 +0000132
133 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
Damien George55baff42014-01-21 21:40:13 +0000134 int alloc_len = lhs_len + rhs_len;
Damien George5fa93b62014-01-22 14:35:10 +0000135
136 /* code for making qstr
Damien George55baff42014-01-21 21:40:13 +0000137 byte *q_ptr;
138 byte *val = qstr_build_start(alloc_len, &q_ptr);
139 memcpy(val, lhs_data, lhs_len);
140 memcpy(val + lhs_len, rhs_data, rhs_len);
Damien George5fa93b62014-01-22 14:35:10 +0000141 return MP_OBJ_NEW_QSTR(qstr_build_end(q_ptr));
142 */
143
144 // code for non-qstr
145 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200146 mp_obj_t s = mp_obj_str_builder_start(mp_obj_get_type(lhs_in), alloc_len, &data);
Damien George5fa93b62014-01-22 14:35:10 +0000147 memcpy(data, lhs_data, lhs_len);
148 memcpy(data + lhs_len, rhs_data, rhs_len);
149 return mp_obj_str_builder_end(s);
Damiend99b0522013-12-21 18:17:45 +0000150 }
151 break;
Damien George5fa93b62014-01-22 14:35:10 +0000152
Damien George9aa2a522014-02-01 23:04:09 +0000153 case RT_BINARY_OP_IN:
John R. Lentonc1bef212014-01-11 12:39:33 +0000154 /* NOTE `a in b` is `b.__contains__(a)` */
Damien George5fa93b62014-01-22 14:35:10 +0000155 if (MP_OBJ_IS_STR(rhs_in)) {
156 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
Damien George9aa2a522014-02-01 23:04:09 +0000157 return MP_BOOL(find_subbytes(lhs_data, lhs_len, rhs_data, rhs_len) != NULL);
John R. Lentonc1bef212014-01-11 12:39:33 +0000158 }
159 break;
Damien George5fa93b62014-01-22 14:35:10 +0000160
Paul Sokolovsky545591a2014-01-21 00:27:33 +0200161 case RT_BINARY_OP_MULTIPLY:
162 {
163 if (!MP_OBJ_IS_SMALL_INT(rhs_in)) {
164 return NULL;
165 }
166 int n = MP_OBJ_SMALL_INT_VALUE(rhs_in);
Damien George5fa93b62014-01-22 14:35:10 +0000167 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200168 mp_obj_t s = mp_obj_str_builder_start(mp_obj_get_type(lhs_in), lhs_len * n, &data);
Damien George5fa93b62014-01-22 14:35:10 +0000169 mp_seq_multiply(lhs_data, sizeof(*lhs_data), lhs_len, n, data);
170 return mp_obj_str_builder_end(s);
Paul Sokolovsky545591a2014-01-21 00:27:33 +0200171 }
Paul Sokolovsky87e85b72014-02-02 08:24:07 +0200172
173 // These 2 are never passed here, dealt with as a special case in rt_binary_op().
174 //case RT_BINARY_OP_EQUAL:
175 //case RT_BINARY_OP_NOT_EQUAL:
176 case RT_BINARY_OP_LESS:
177 case RT_BINARY_OP_LESS_EQUAL:
178 case RT_BINARY_OP_MORE:
179 case RT_BINARY_OP_MORE_EQUAL:
180 if (MP_OBJ_IS_STR(rhs_in)) {
181 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
182 return MP_BOOL(mp_seq_cmp_bytes(op, lhs_data, lhs_len, rhs_data, rhs_len));
183 }
Damiend99b0522013-12-21 18:17:45 +0000184 }
185
186 return MP_OBJ_NULL; // op not supported
187}
188
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200189STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
Damien George5fa93b62014-01-22 14:35:10 +0000190 assert(MP_OBJ_IS_STR(self_in));
Damiend99b0522013-12-21 18:17:45 +0000191
Damien Georgefe8fb912014-01-02 16:36:09 +0000192 // get separation string
Damien George5fa93b62014-01-22 14:35:10 +0000193 GET_STR_DATA_LEN(self_in, sep_str, sep_len);
Damien Georgefe8fb912014-01-02 16:36:09 +0000194
195 // process args
Damiend99b0522013-12-21 18:17:45 +0000196 uint seq_len;
197 mp_obj_t *seq_items;
198 if (MP_OBJ_IS_TYPE(arg, &tuple_type)) {
199 mp_obj_tuple_get(arg, &seq_len, &seq_items);
200 } else if (MP_OBJ_IS_TYPE(arg, &list_type)) {
201 mp_obj_list_get(arg, &seq_len, &seq_items);
202 } else {
203 goto bad_arg;
204 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000205
206 // count required length
207 int required_len = 0;
Damiend99b0522013-12-21 18:17:45 +0000208 for (int i = 0; i < seq_len; i++) {
Damien George5fa93b62014-01-22 14:35:10 +0000209 if (!MP_OBJ_IS_STR(seq_items[i])) {
Damiend99b0522013-12-21 18:17:45 +0000210 goto bad_arg;
211 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000212 if (i > 0) {
213 required_len += sep_len;
214 }
Damien George5fa93b62014-01-22 14:35:10 +0000215 GET_STR_LEN(seq_items[i], l);
216 required_len += l;
Damiend99b0522013-12-21 18:17:45 +0000217 }
218
219 // make joined string
Damien George5fa93b62014-01-22 14:35:10 +0000220 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200221 mp_obj_t joined_str = mp_obj_str_builder_start(mp_obj_get_type(self_in), required_len, &data);
Damiend99b0522013-12-21 18:17:45 +0000222 for (int i = 0; i < seq_len; i++) {
Damiend99b0522013-12-21 18:17:45 +0000223 if (i > 0) {
Damien George5fa93b62014-01-22 14:35:10 +0000224 memcpy(data, sep_str, sep_len);
225 data += sep_len;
Damiend99b0522013-12-21 18:17:45 +0000226 }
Damien George5fa93b62014-01-22 14:35:10 +0000227 GET_STR_DATA_LEN(seq_items[i], s, l);
228 memcpy(data, s, l);
229 data += l;
Damiend99b0522013-12-21 18:17:45 +0000230 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000231
232 // return joined string
Damien George5fa93b62014-01-22 14:35:10 +0000233 return mp_obj_str_builder_end(joined_str);
Damiend99b0522013-12-21 18:17:45 +0000234
235bad_arg:
Damien Georgec5966122014-02-15 16:10:44 +0000236 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "?str.join expecting a list of str's"));
Damiend99b0522013-12-21 18:17:45 +0000237}
238
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200239#define is_ws(c) ((c) == ' ' || (c) == '\t')
240
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200241STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200242 int splits = -1;
243 mp_obj_t sep = mp_const_none;
244 if (n_args > 1) {
245 sep = args[1];
246 if (n_args > 2) {
247 splits = MP_OBJ_SMALL_INT_VALUE(args[2]);
248 }
249 }
250 assert(sep == mp_const_none);
Damien George12eacca2014-01-21 21:54:15 +0000251 (void)sep; // unused; to hush compiler warning
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200252 mp_obj_t res = mp_obj_new_list(0, NULL);
Damien George5fa93b62014-01-22 14:35:10 +0000253 GET_STR_DATA_LEN(args[0], s, len);
254 const byte *top = s + len;
255 const byte *start;
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200256
257 // Initial whitespace is not counted as split, so we pre-do it
Damien George5fa93b62014-01-22 14:35:10 +0000258 while (s < top && is_ws(*s)) s++;
259 while (s < top && splits != 0) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200260 start = s;
Damien George5fa93b62014-01-22 14:35:10 +0000261 while (s < top && !is_ws(*s)) s++;
262 rt_list_append(res, mp_obj_new_str(start, s - start, false));
263 if (s >= top) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200264 break;
265 }
Damien George5fa93b62014-01-22 14:35:10 +0000266 while (s < top && is_ws(*s)) s++;
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200267 if (splits > 0) {
268 splits--;
269 }
270 }
271
Damien George5fa93b62014-01-22 14:35:10 +0000272 if (s < top) {
273 rt_list_append(res, mp_obj_new_str(s, top - s, false));
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200274 }
275
276 return res;
277}
278
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200279STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
John R. Lentone8204912014-01-12 21:53:52 +0000280 assert(2 <= n_args && n_args <= 4);
Damien George5fa93b62014-01-22 14:35:10 +0000281 assert(MP_OBJ_IS_STR(args[0]));
282 assert(MP_OBJ_IS_STR(args[1]));
John R. Lentone8204912014-01-12 21:53:52 +0000283
Damien George5fa93b62014-01-22 14:35:10 +0000284 GET_STR_DATA_LEN(args[0], haystack, haystack_len);
285 GET_STR_DATA_LEN(args[1], needle, needle_len);
John R. Lentone8204912014-01-12 21:53:52 +0000286
xbec5538882014-03-16 17:58:35 -0700287 machine_uint_t start = 0;
288 machine_uint_t end = haystack_len;
John R. Lentone8204912014-01-12 21:53:52 +0000289 /* TODO use a non-exception-throwing mp_get_index */
290 if (n_args >= 3 && args[2] != mp_const_none) {
xbe9e1e8cd2014-03-12 22:57:16 -0700291 start = mp_get_index(&str_type, haystack_len, args[2], true);
John R. Lentone8204912014-01-12 21:53:52 +0000292 }
293 if (n_args >= 4 && args[3] != mp_const_none) {
xbe9e1e8cd2014-03-12 22:57:16 -0700294 end = mp_get_index(&str_type, haystack_len, args[3], true);
John R. Lentone8204912014-01-12 21:53:52 +0000295 }
296
Damien George5fa93b62014-01-22 14:35:10 +0000297 const byte *p = find_subbytes(haystack + start, haystack_len - start, needle, needle_len);
Damien George23005372014-01-13 19:39:01 +0000298 if (p == NULL) {
299 // not found
300 return MP_OBJ_NEW_SMALL_INT(-1);
301 } else {
302 // found
303 machine_int_t pos = p - haystack;
John R. Lentone8204912014-01-12 21:53:52 +0000304 if (pos + needle_len > end) {
305 pos = -1;
306 }
Damien George23005372014-01-13 19:39:01 +0000307 return MP_OBJ_NEW_SMALL_INT(pos);
John R. Lentone8204912014-01-12 21:53:52 +0000308 }
John R. Lentone8204912014-01-12 21:53:52 +0000309}
310
Paul Sokolovsky1eacefe2014-01-23 01:20:40 +0200311// TODO: (Much) more variety in args
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200312STATIC mp_obj_t str_startswith(mp_obj_t self_in, mp_obj_t arg) {
Paul Sokolovsky1eacefe2014-01-23 01:20:40 +0200313 GET_STR_DATA_LEN(self_in, str, str_len);
314 GET_STR_DATA_LEN(arg, prefix, prefix_len);
315 if (prefix_len > str_len) {
316 return mp_const_false;
317 }
318 return MP_BOOL(memcmp(str, prefix, prefix_len) == 0);
319}
320
xbec5538882014-03-16 17:58:35 -0700321STATIC bool chr_in_str(const byte* const str, const machine_uint_t str_len, int c) {
322 for (machine_uint_t i = 0; i < str_len; i++) {
Damien George5fa93b62014-01-22 14:35:10 +0000323 if (str[i] == c) {
324 return true;
325 }
326 }
327 return false;
328}
329
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200330STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) {
xbe7b0f39f2014-01-08 14:23:45 -0800331 assert(1 <= n_args && n_args <= 2);
Damien George5fa93b62014-01-22 14:35:10 +0000332 assert(MP_OBJ_IS_STR(args[0]));
333
334 const byte *chars_to_del;
335 uint chars_to_del_len;
336 static const byte whitespace[] = " \t\n\r\v\f";
xbe7b0f39f2014-01-08 14:23:45 -0800337
338 if (n_args == 1) {
339 chars_to_del = whitespace;
Damien George5fa93b62014-01-22 14:35:10 +0000340 chars_to_del_len = sizeof(whitespace);
xbe7b0f39f2014-01-08 14:23:45 -0800341 } else {
Damien George5fa93b62014-01-22 14:35:10 +0000342 assert(MP_OBJ_IS_STR(args[1]));
343 GET_STR_DATA_LEN(args[1], s, l);
344 chars_to_del = s;
345 chars_to_del_len = l;
xbe7b0f39f2014-01-08 14:23:45 -0800346 }
347
Damien George5fa93b62014-01-22 14:35:10 +0000348 GET_STR_DATA_LEN(args[0], orig_str, orig_str_len);
xbe7b0f39f2014-01-08 14:23:45 -0800349
xbec5538882014-03-16 17:58:35 -0700350 machine_uint_t first_good_char_pos = 0;
xbe7b0f39f2014-01-08 14:23:45 -0800351 bool first_good_char_pos_set = false;
xbec5538882014-03-16 17:58:35 -0700352 machine_uint_t last_good_char_pos = 0;
353 for (machine_uint_t i = 0; i < orig_str_len; i++) {
xbe7b0f39f2014-01-08 14:23:45 -0800354 if (!chr_in_str(chars_to_del, chars_to_del_len, orig_str[i])) {
355 last_good_char_pos = i;
356 if (!first_good_char_pos_set) {
357 first_good_char_pos = i;
358 first_good_char_pos_set = true;
359 }
360 }
361 }
362
363 if (first_good_char_pos == 0 && last_good_char_pos == 0) {
Damien George5fa93b62014-01-22 14:35:10 +0000364 // string is all whitespace, return ''
365 return MP_OBJ_NEW_QSTR(MP_QSTR_);
xbe7b0f39f2014-01-08 14:23:45 -0800366 }
367
368 assert(last_good_char_pos >= first_good_char_pos);
369 //+1 to accomodate the last character
xbec5538882014-03-16 17:58:35 -0700370 machine_uint_t stripped_len = last_good_char_pos - first_good_char_pos + 1;
Damien George5fa93b62014-01-22 14:35:10 +0000371 return mp_obj_new_str(orig_str + first_good_char_pos, stripped_len, false);
xbe7b0f39f2014-01-08 14:23:45 -0800372}
373
Damien Georgea11ceca2014-01-19 16:02:09 +0000374mp_obj_t str_format(uint n_args, const mp_obj_t *args) {
Damien George5fa93b62014-01-22 14:35:10 +0000375 assert(MP_OBJ_IS_STR(args[0]));
Damiend99b0522013-12-21 18:17:45 +0000376
Damien George5fa93b62014-01-22 14:35:10 +0000377 GET_STR_DATA_LEN(args[0], str, len);
Damiend99b0522013-12-21 18:17:45 +0000378 int arg_i = 1;
379 vstr_t *vstr = vstr_new();
Damien George5fa93b62014-01-22 14:35:10 +0000380 for (const byte *top = str + len; str < top; str++) {
Damiend99b0522013-12-21 18:17:45 +0000381 if (*str == '{') {
382 str++;
Damien George5fa93b62014-01-22 14:35:10 +0000383 if (str < top && *str == '{') {
Damiend99b0522013-12-21 18:17:45 +0000384 vstr_add_char(vstr, '{');
Paul Sokolovskyf2b796e2014-01-15 22:45:20 +0200385 } else {
Damien George5fa93b62014-01-22 14:35:10 +0000386 while (str < top && *str != '}') str++;
Damiend99b0522013-12-21 18:17:45 +0000387 if (arg_i >= n_args) {
Damien Georgec5966122014-02-15 16:10:44 +0000388 nlr_jump(mp_obj_new_exception_msg(&mp_type_IndexError, "tuple index out of range"));
Damiend99b0522013-12-21 18:17:45 +0000389 }
Paul Sokolovsky76d982e2014-01-13 19:19:16 +0200390 // TODO: may be PRINT_REPR depending on formatting code
Damien George4899ff92014-01-15 22:39:03 +0000391 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR);
Damiend99b0522013-12-21 18:17:45 +0000392 arg_i++;
393 }
394 } else {
395 vstr_add_char(vstr, *str);
396 }
397 }
398
Damien George5fa93b62014-01-22 14:35:10 +0000399 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
400 vstr_free(vstr);
401 return s;
Damiend99b0522013-12-21 18:17:45 +0000402}
403
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200404STATIC mp_obj_t str_replace(uint n_args, const mp_obj_t *args) {
xbe480c15a2014-01-30 22:17:30 -0800405 assert(MP_OBJ_IS_STR(args[0]));
406 assert(MP_OBJ_IS_STR(args[1]));
407 assert(MP_OBJ_IS_STR(args[2]));
408
Damien George94f68302014-01-31 23:45:12 +0000409 machine_int_t max_rep = 0;
xbe480c15a2014-01-30 22:17:30 -0800410 if (n_args == 4) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200411 assert(MP_OBJ_IS_SMALL_INT(args[3]));
412 max_rep = MP_OBJ_SMALL_INT_VALUE(args[3]);
413 if (max_rep == 0) {
414 return args[0];
415 } else if (max_rep < 0) {
416 max_rep = 0;
417 }
xbe480c15a2014-01-30 22:17:30 -0800418 }
Damien George94f68302014-01-31 23:45:12 +0000419
420 // if max_rep is still 0 by this point we will need to do all possible replacements
xbe480c15a2014-01-30 22:17:30 -0800421
422 GET_STR_DATA_LEN(args[0], str, str_len);
423 GET_STR_DATA_LEN(args[1], old, old_len);
424 GET_STR_DATA_LEN(args[2], new, new_len);
Damien George94f68302014-01-31 23:45:12 +0000425
426 // old won't exist in str if it's longer, so nothing to replace
xbe480c15a2014-01-30 22:17:30 -0800427 if (old_len > str_len) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200428 return args[0];
xbe480c15a2014-01-30 22:17:30 -0800429 }
430
Damien George94f68302014-01-31 23:45:12 +0000431 // data for the replaced string
432 byte *data = NULL;
433 mp_obj_t replaced_str = MP_OBJ_NULL;
xbe480c15a2014-01-30 22:17:30 -0800434
Damien George94f68302014-01-31 23:45:12 +0000435 // do 2 passes over the string:
436 // first pass computes the required length of the replaced string
437 // second pass does the replacements
438 for (;;) {
439 machine_uint_t replaced_str_index = 0;
440 machine_uint_t num_replacements_done = 0;
441 const byte *old_occurrence;
442 const byte *offset_ptr = str;
443 machine_uint_t offset_num = 0;
444 while ((old_occurrence = find_subbytes(offset_ptr, str_len - offset_num, old, old_len)) != NULL) {
445 // copy from just after end of last occurrence of to-be-replaced string to right before start of next occurrence
446 if (data != NULL) {
447 memcpy(data + replaced_str_index, offset_ptr, old_occurrence - offset_ptr);
448 }
449 replaced_str_index += old_occurrence - offset_ptr;
450 // copy the replacement string
451 if (data != NULL) {
452 memcpy(data + replaced_str_index, new, new_len);
453 }
454 replaced_str_index += new_len;
455 offset_ptr = old_occurrence + old_len;
456 offset_num = offset_ptr - str;
xbe480c15a2014-01-30 22:17:30 -0800457
Damien George94f68302014-01-31 23:45:12 +0000458 num_replacements_done++;
459 if (max_rep != 0 && num_replacements_done == max_rep){
460 break;
461 }
462 }
463
464 // copy from just after end of last occurrence of to-be-replaced string to end of old string
465 if (data != NULL) {
466 memcpy(data + replaced_str_index, offset_ptr, str_len - offset_num);
467 }
468 replaced_str_index += str_len - offset_num;
469
470 if (data == NULL) {
471 // first pass
472 if (num_replacements_done == 0) {
473 // no substr found, return original string
474 return args[0];
475 } else {
476 // substr found, allocate new string
477 replaced_str = mp_obj_str_builder_start(mp_obj_get_type(args[0]), replaced_str_index, &data);
478 }
479 } else {
480 // second pass, we are done
481 break;
482 }
xbe480c15a2014-01-30 22:17:30 -0800483 }
Damien George94f68302014-01-31 23:45:12 +0000484
xbe480c15a2014-01-30 22:17:30 -0800485 return mp_obj_str_builder_end(replaced_str);
486}
487
xbe9e1e8cd2014-03-12 22:57:16 -0700488STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
489 assert(2 <= n_args && n_args <= 4);
490 assert(MP_OBJ_IS_STR(args[0]));
491 assert(MP_OBJ_IS_STR(args[1]));
492
493 GET_STR_DATA_LEN(args[0], haystack, haystack_len);
494 GET_STR_DATA_LEN(args[1], needle, needle_len);
495
Damien George536dde22014-03-13 22:07:55 +0000496 machine_uint_t start = 0;
497 machine_uint_t end = haystack_len;
xbe9e1e8cd2014-03-12 22:57:16 -0700498 /* TODO use a non-exception-throwing mp_get_index */
499 if (n_args >= 3 && args[2] != mp_const_none) {
500 start = mp_get_index(&str_type, haystack_len, args[2], true);
501 }
502 if (n_args >= 4 && args[3] != mp_const_none) {
503 end = mp_get_index(&str_type, haystack_len, args[3], true);
504 }
505
Damien George536dde22014-03-13 22:07:55 +0000506 // if needle_len is zero then we count each gap between characters as an occurrence
507 if (needle_len == 0) {
508 return MP_OBJ_NEW_SMALL_INT(end - start + 1);
xbe9e1e8cd2014-03-12 22:57:16 -0700509 }
510
Damien George536dde22014-03-13 22:07:55 +0000511 // count the occurrences
512 machine_int_t num_occurrences = 0;
xbec5d70ba2014-03-13 00:29:15 -0700513 for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
514 if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
515 num_occurrences++;
516 haystack_index += needle_len - 1;
517 }
xbe9e1e8cd2014-03-12 22:57:16 -0700518 }
519
520 return MP_OBJ_NEW_SMALL_INT(num_occurrences);
521}
522
xbe613a8e32014-03-18 00:06:29 -0700523STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
524 assert(MP_OBJ_IS_STR(self_in));
525 if (!MP_OBJ_IS_STR(arg)) {
526 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
527 "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
528 }
529
530 GET_STR_DATA_LEN(self_in, str, str_len);
531 GET_STR_DATA_LEN(arg, sep, sep_len);
532
533 if (sep_len == 0) {
534 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
535 }
536
537 for (machine_uint_t str_index = 0; str_index + sep_len <= str_len; str_index++) {
538 if (memcmp(&str[str_index], sep, sep_len) == 0) {
539 mp_obj_t items[] = {mp_obj_new_str(str, str_index, false), arg,
540 mp_obj_new_str(str + str_index + sep_len, str_len - str_index - sep_len, false)};
541 return mp_obj_new_tuple(3, items);
542 }
543 }
544 mp_obj_t items[] = {mp_obj_new_str(str, str_len, false), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
545 return mp_obj_new_tuple(3, items);
546}
547
xbe4504ea82014-03-19 00:46:14 -0700548STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg) {
549 assert(MP_OBJ_IS_STR(self_in));
550 if (!MP_OBJ_IS_STR(arg)) {
551 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
552 "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
553 }
554
555 GET_STR_DATA_LEN(self_in, str, str_len);
556 GET_STR_DATA_LEN(arg, sep, sep_len);
557
558 if (sep_len == 0) {
559 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
560 }
561
562 if (sep_len > str_len) {
563 goto not_found;
564 }
565
566 for (machine_uint_t str_index = str_len; ; str_index--) {
567 if (memcmp(&str[str_index - sep_len], sep, sep_len) == 0) {
568 mp_obj_t items[] = {mp_obj_new_str(str, str_index - sep_len, false), arg,
569 mp_obj_new_str(str + str_index, str_len - str_index, false)};
570 return mp_obj_new_tuple(3, items);
571 }
572 if (str_index - sep_len == 0) {
573 break;
574 }
575 }
576
577not_found: ;
578 mp_obj_t items[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), mp_obj_new_str(str, str_len, false)};
579 return mp_obj_new_tuple(3, items);
580}
581
Damien George2da98302014-03-09 19:58:18 +0000582STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
583 if (flags == BUFFER_READ) {
584 GET_STR_DATA_LEN(self_in, str_data, str_len);
585 bufinfo->buf = (void*)str_data;
586 bufinfo->len = str_len;
587 return 0;
588 } else {
589 // can't write to a string
590 bufinfo->buf = NULL;
591 bufinfo->len = 0;
592 return 1;
593 }
594}
595
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200596STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
597STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
598STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
599STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
600STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
601STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
602STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
xbe9e1e8cd2014-03-12 22:57:16 -0700603STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
xbe613a8e32014-03-18 00:06:29 -0700604STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_partition_obj, str_partition);
xbe4504ea82014-03-19 00:46:14 -0700605STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition);
Damiend99b0522013-12-21 18:17:45 +0000606
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200607STATIC const mp_method_t str_type_methods[] = {
John R. Lentone8204912014-01-12 21:53:52 +0000608 { "find", &str_find_obj },
ian-v7a16fad2014-01-06 09:52:29 -0800609 { "join", &str_join_obj },
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200610 { "split", &str_split_obj },
Paul Sokolovsky1eacefe2014-01-23 01:20:40 +0200611 { "startswith", &str_startswith_obj },
xbe7b0f39f2014-01-08 14:23:45 -0800612 { "strip", &str_strip_obj },
ian-v7a16fad2014-01-06 09:52:29 -0800613 { "format", &str_format_obj },
xbe480c15a2014-01-30 22:17:30 -0800614 { "replace", &str_replace_obj },
xbe9e1e8cd2014-03-12 22:57:16 -0700615 { "count", &str_count_obj },
xbe613a8e32014-03-18 00:06:29 -0700616 { "partition", &str_partition_obj },
xbe4504ea82014-03-19 00:46:14 -0700617 { "rpartition", &str_rpartition_obj },
ian-v7a16fad2014-01-06 09:52:29 -0800618 { NULL, NULL }, // end-of-list sentinel
619};
Damien George97209d32014-01-07 15:58:30 +0000620
Damiend99b0522013-12-21 18:17:45 +0000621const mp_obj_type_t str_type = {
Damien Georgec5966122014-02-15 16:10:44 +0000622 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000623 .name = MP_QSTR_str,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +0200624 .print = str_print,
625 .binary_op = str_binary_op,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200626 .getiter = mp_obj_new_str_iterator,
ian-v7a16fad2014-01-06 09:52:29 -0800627 .methods = str_type_methods,
Damien George2da98302014-03-09 19:58:18 +0000628 .buffer_p = { .get_buffer = str_get_buffer },
Damiend99b0522013-12-21 18:17:45 +0000629};
630
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200631// Reuses most of methods from str
632const mp_obj_type_t bytes_type = {
Damien Georgec5966122014-02-15 16:10:44 +0000633 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000634 .name = MP_QSTR_bytes,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200635 .print = str_print,
636 .binary_op = str_binary_op,
637 .getiter = mp_obj_new_bytes_iterator,
638 .methods = str_type_methods,
639};
640
641mp_obj_t mp_obj_str_builder_start(const mp_obj_type_t *type, uint len, byte **data) {
Damien George5fa93b62014-01-22 14:35:10 +0000642 mp_obj_str_t *o = m_new_obj_var(mp_obj_str_t, byte, len + 1);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200643 o->base.type = type;
Damien George5fa93b62014-01-22 14:35:10 +0000644 o->len = len;
645 *data = o->data;
Damiend99b0522013-12-21 18:17:45 +0000646 return o;
647}
648
Damien George5fa93b62014-01-22 14:35:10 +0000649mp_obj_t mp_obj_str_builder_end(mp_obj_t o_in) {
650 assert(MP_OBJ_IS_STR(o_in));
651 mp_obj_str_t *o = o_in;
652 o->hash = qstr_compute_hash(o->data, o->len);
653 o->data[o->len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
654 return o;
655}
656
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200657STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len) {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200658 mp_obj_str_t *o = m_new_obj_var(mp_obj_str_t, byte, len + 1);
659 o->base.type = type;
660 o->hash = qstr_compute_hash(data, len);
661 o->len = len;
662 memcpy(o->data, data, len * sizeof(byte));
663 o->data[len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
664 return o;
665}
666
Damien George5fa93b62014-01-22 14:35:10 +0000667mp_obj_t mp_obj_new_str(const byte* data, uint len, bool make_qstr_if_not_already) {
668 qstr q = qstr_find_strn(data, len);
669 if (q != MP_QSTR_NULL) {
670 // qstr with this data already exists
671 return MP_OBJ_NEW_QSTR(q);
672 } else if (make_qstr_if_not_already) {
673 // no existing qstr, make a new one
674 return MP_OBJ_NEW_QSTR(qstr_from_strn((const char*)data, len));
675 } else {
676 // no existing qstr, don't make one
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200677 return str_new(&str_type, data, len);
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +0200678 }
Damien George5fa93b62014-01-22 14:35:10 +0000679}
680
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200681mp_obj_t mp_obj_new_bytes(const byte* data, uint len) {
682 return str_new(&bytes_type, data, len);
683}
684
Damien George5fa93b62014-01-22 14:35:10 +0000685bool mp_obj_str_equal(mp_obj_t s1, mp_obj_t s2) {
686 if (MP_OBJ_IS_QSTR(s1) && MP_OBJ_IS_QSTR(s2)) {
687 return s1 == s2;
688 } else {
689 GET_STR_HASH(s1, h1);
690 GET_STR_HASH(s2, h2);
691 if (h1 != h2) {
692 return false;
693 }
694 GET_STR_DATA_LEN(s1, d1, l1);
695 GET_STR_DATA_LEN(s2, d2, l2);
696 if (l1 != l2) {
697 return false;
698 }
Damien George1e708fe2014-01-23 18:27:51 +0000699 return memcmp(d1, d2, l1) == 0;
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +0200700 }
Damien George5fa93b62014-01-22 14:35:10 +0000701}
702
Damien Georgeb829b5c2014-01-25 13:51:19 +0000703void bad_implicit_conversion(mp_obj_t self_in) __attribute__((noreturn));
704void bad_implicit_conversion(mp_obj_t self_in) {
Damien Georgec5966122014-02-15 16:10:44 +0000705 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(self_in)));
Damien Georgeb829b5c2014-01-25 13:51:19 +0000706}
707
Damien George5fa93b62014-01-22 14:35:10 +0000708uint mp_obj_str_get_hash(mp_obj_t self_in) {
709 if (MP_OBJ_IS_STR(self_in)) {
710 GET_STR_HASH(self_in, h);
711 return h;
712 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +0000713 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +0000714 }
715}
716
717uint mp_obj_str_get_len(mp_obj_t self_in) {
718 if (MP_OBJ_IS_STR(self_in)) {
719 GET_STR_LEN(self_in, l);
720 return l;
721 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +0000722 bad_implicit_conversion(self_in);
723 }
724}
725
726// use this if you will anyway convert the string to a qstr
727// will be more efficient for the case where it's already a qstr
728qstr mp_obj_str_get_qstr(mp_obj_t self_in) {
729 if (MP_OBJ_IS_QSTR(self_in)) {
730 return MP_OBJ_QSTR_VALUE(self_in);
731 } else if (MP_OBJ_IS_TYPE(self_in, &str_type)) {
732 mp_obj_str_t *self = self_in;
733 return qstr_from_strn((char*)self->data, self->len);
734 } else {
735 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +0000736 }
737}
738
739// only use this function if you need the str data to be zero terminated
740// at the moment all strings are zero terminated to help with C ASCIIZ compatibility
741const char *mp_obj_str_get_str(mp_obj_t self_in) {
742 if (MP_OBJ_IS_STR(self_in)) {
743 GET_STR_DATA_LEN(self_in, s, l);
744 (void)l; // len unused
745 return (const char*)s;
746 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +0000747 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +0000748 }
749}
750
Damien George698ec212014-02-08 18:17:23 +0000751const char *mp_obj_str_get_data(mp_obj_t self_in, uint *len) {
Damien George5fa93b62014-01-22 14:35:10 +0000752 if (MP_OBJ_IS_STR(self_in)) {
753 GET_STR_DATA_LEN(self_in, s, l);
754 *len = l;
Damien George698ec212014-02-08 18:17:23 +0000755 return (const char*)s;
Damien George5fa93b62014-01-22 14:35:10 +0000756 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +0000757 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +0000758 }
Damiend99b0522013-12-21 18:17:45 +0000759}
xyb8cfc9f02014-01-05 18:47:51 +0800760
761/******************************************************************************/
762/* str iterator */
763
764typedef struct _mp_obj_str_it_t {
765 mp_obj_base_t base;
Damien George5fa93b62014-01-22 14:35:10 +0000766 mp_obj_t str;
xyb8cfc9f02014-01-05 18:47:51 +0800767 machine_uint_t cur;
768} mp_obj_str_it_t;
769
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200770STATIC mp_obj_t str_it_iternext(mp_obj_t self_in) {
xyb8cfc9f02014-01-05 18:47:51 +0800771 mp_obj_str_it_t *self = self_in;
Damien George5fa93b62014-01-22 14:35:10 +0000772 GET_STR_DATA_LEN(self->str, str, len);
773 if (self->cur < len) {
774 mp_obj_t o_out = mp_obj_new_str(str + self->cur, 1, true);
xyb8cfc9f02014-01-05 18:47:51 +0800775 self->cur += 1;
776 return o_out;
777 } else {
778 return mp_const_stop_iteration;
779 }
780}
781
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200782STATIC const mp_obj_type_t str_it_type = {
Damien Georgec5966122014-02-15 16:10:44 +0000783 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000784 .name = MP_QSTR_iterator,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +0200785 .iternext = str_it_iternext,
xyb8cfc9f02014-01-05 18:47:51 +0800786};
787
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200788STATIC mp_obj_t bytes_it_iternext(mp_obj_t self_in) {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200789 mp_obj_str_it_t *self = self_in;
790 GET_STR_DATA_LEN(self->str, str, len);
791 if (self->cur < len) {
Damien George7c9c6672014-01-25 00:17:36 +0000792 mp_obj_t o_out = MP_OBJ_NEW_SMALL_INT((mp_small_int_t)str[self->cur]);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200793 self->cur += 1;
794 return o_out;
795 } else {
796 return mp_const_stop_iteration;
797 }
798}
799
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200800STATIC const mp_obj_type_t bytes_it_type = {
Damien Georgec5966122014-02-15 16:10:44 +0000801 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000802 .name = MP_QSTR_iterator,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200803 .iternext = bytes_it_iternext,
804};
805
806mp_obj_t mp_obj_new_str_iterator(mp_obj_t str) {
xyb8cfc9f02014-01-05 18:47:51 +0800807 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
808 o->base.type = &str_it_type;
809 o->str = str;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200810 o->cur = 0;
811 return o;
812}
813
814mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str) {
815 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
816 o->base.type = &bytes_it_type;
817 o->str = str;
818 o->cur = 0;
xyb8cfc9f02014-01-05 18:47:51 +0800819 return o;
820}