blob: 8389bb0bdfb3043ec2fff88ba92a46e189a07ddb [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"
Dave Hylandsbaf6f142014-03-30 21:06:50 -070012#include "pfenv.h"
Damiend99b0522013-12-21 18:17:45 +000013
14typedef struct _mp_obj_str_t {
15 mp_obj_base_t base;
Damien George5fa93b62014-01-22 14:35:10 +000016 machine_uint_t hash : 16; // XXX here we assume the hash size is 16 bits (it is at the moment; see qstr.c)
17 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
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +020018 const byte *data;
Damiend99b0522013-12-21 18:17:45 +000019} mp_obj_str_t;
20
Paul Sokolovsky4db727a2014-03-31 21:18:28 +030021STATIC mp_obj_t str_modulo_format(mp_obj_t pattern, uint n_args, const mp_obj_t *args);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +020022const mp_obj_t mp_const_empty_bytes;
23
Damien George5fa93b62014-01-22 14:35:10 +000024// use this macro to extract the string hash
25#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; }
26
27// use this macro to extract the string length
28#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; }
29
30// use this macro to extract the string data and length
31#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; }
32
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020033STATIC mp_obj_t mp_obj_new_str_iterator(mp_obj_t str);
34STATIC mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str);
Paul Sokolovskybe020c22014-03-21 11:39:01 +020035STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len);
xyb8cfc9f02014-01-05 18:47:51 +080036
37/******************************************************************************/
38/* str */
39
Paul Sokolovsky0b7e29c2014-01-28 03:40:06 +020040void mp_str_print_quoted(void (*print)(void *env, const char *fmt, ...), void *env, const byte *str_data, uint str_len) {
41 // this escapes characters, but it will be very slow to print (calling print many times)
42 bool has_single_quote = false;
43 bool has_double_quote = false;
44 for (const byte *s = str_data, *top = str_data + str_len; (!has_single_quote || !has_double_quote) && s < top; s++) {
45 if (*s == '\'') {
46 has_single_quote = true;
47 } else if (*s == '"') {
48 has_double_quote = true;
49 }
50 }
51 int quote_char = '\'';
52 if (has_single_quote && !has_double_quote) {
53 quote_char = '"';
54 }
55 print(env, "%c", quote_char);
56 for (const byte *s = str_data, *top = str_data + str_len; s < top; s++) {
57 if (*s == quote_char) {
58 print(env, "\\%c", quote_char);
59 } else if (*s == '\\') {
60 print(env, "\\\\");
61 } else if (32 <= *s && *s <= 126) {
62 print(env, "%c", *s);
63 } else if (*s == '\n') {
64 print(env, "\\n");
65 // TODO add more escape codes here if we want to match CPython
66 } else {
67 print(env, "\\x%02x", *s);
68 }
69 }
70 print(env, "%c", quote_char);
71}
72
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +020073STATIC 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 +000074 GET_STR_DATA_LEN(self_in, str_data, str_len);
Damien George3e1a5c12014-03-29 13:43:38 +000075 bool is_bytes = MP_OBJ_IS_TYPE(self_in, &mp_type_bytes);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +020076 if (kind == PRINT_STR && !is_bytes) {
Damien George5fa93b62014-01-22 14:35:10 +000077 print(env, "%.*s", str_len, str_data);
Paul Sokolovsky76d982e2014-01-13 19:19:16 +020078 } else {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +020079 if (is_bytes) {
80 print(env, "b");
81 }
Paul Sokolovsky0b7e29c2014-01-28 03:40:06 +020082 mp_str_print_quoted(print, env, str_data, str_len);
Paul Sokolovsky76d982e2014-01-13 19:19:16 +020083 }
Damiend99b0522013-12-21 18:17:45 +000084}
85
Paul Sokolovskybe020c22014-03-21 11:39:01 +020086STATIC mp_obj_t str_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
87 switch (n_args) {
88 case 0:
89 return MP_OBJ_NEW_QSTR(MP_QSTR_);
90
91 case 1:
92 {
93 vstr_t *vstr = vstr_new();
94 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[0], PRINT_STR);
95 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
96 vstr_free(vstr);
97 return s;
98 }
99
100 case 2:
101 case 3:
102 {
103 // TODO: validate 2nd/3rd args
Damien George3e1a5c12014-03-29 13:43:38 +0000104 if (!MP_OBJ_IS_TYPE(args[0], &mp_type_bytes)) {
Paul Sokolovskybe020c22014-03-21 11:39:01 +0200105 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "bytes expected"));
106 }
107 GET_STR_DATA_LEN(args[0], str_data, str_len);
108 GET_STR_HASH(args[0], str_hash);
Damien George3e1a5c12014-03-29 13:43:38 +0000109 mp_obj_str_t *o = str_new(&mp_type_str, NULL, str_len);
Paul Sokolovskybe020c22014-03-21 11:39:01 +0200110 o->data = str_data;
111 o->hash = str_hash;
112 return o;
113 }
114
115 default:
116 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "str takes at most 3 arguments"));
117 }
118}
119
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200120STATIC mp_obj_t bytes_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) {
121 if (n_args == 0) {
122 return mp_const_empty_bytes;
123 }
124
125 if (MP_OBJ_IS_STR(args[0])) {
126 if (n_args < 2 || n_args > 3) {
127 goto wrong_args;
128 }
129 GET_STR_DATA_LEN(args[0], str_data, str_len);
130 GET_STR_HASH(args[0], str_hash);
Damien George3e1a5c12014-03-29 13:43:38 +0000131 mp_obj_str_t *o = str_new(&mp_type_bytes, NULL, str_len);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200132 o->data = str_data;
133 o->hash = str_hash;
134 return o;
135 }
136
137 if (n_args > 1) {
138 goto wrong_args;
139 }
140
141 if (MP_OBJ_IS_SMALL_INT(args[0])) {
142 uint len = MP_OBJ_SMALL_INT_VALUE(args[0]);
143 byte *data;
144
Damien George3e1a5c12014-03-29 13:43:38 +0000145 mp_obj_t o = mp_obj_str_builder_start(&mp_type_bytes, len, &data);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200146 memset(data, 0, len);
147 return mp_obj_str_builder_end(o);
148 }
149
150 int len;
151 byte *data;
152 vstr_t *vstr = NULL;
153 mp_obj_t o = NULL;
154 // Try to create array of exact len if initializer len is known
155 mp_obj_t len_in = mp_obj_len_maybe(args[0]);
156 if (len_in == MP_OBJ_NULL) {
157 len = -1;
158 vstr = vstr_new();
159 } else {
160 len = MP_OBJ_SMALL_INT_VALUE(len_in);
Damien George3e1a5c12014-03-29 13:43:38 +0000161 o = mp_obj_str_builder_start(&mp_type_bytes, len, &data);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200162 }
163
Damien Georged17926d2014-03-30 13:35:08 +0100164 mp_obj_t iterable = mp_getiter(args[0]);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200165 mp_obj_t item;
Damien Georged17926d2014-03-30 13:35:08 +0100166 while ((item = mp_iternext(iterable)) != MP_OBJ_NULL) {
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200167 if (len == -1) {
168 vstr_add_char(vstr, MP_OBJ_SMALL_INT_VALUE(item));
169 } else {
170 *data++ = MP_OBJ_SMALL_INT_VALUE(item);
171 }
172 }
173
174 if (len == -1) {
175 vstr_shrink(vstr);
176 // TODO: Optimize, borrow buffer from vstr
177 len = vstr_len(vstr);
Damien George3e1a5c12014-03-29 13:43:38 +0000178 o = mp_obj_str_builder_start(&mp_type_bytes, len, &data);
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +0200179 memcpy(data, vstr_str(vstr), len);
180 vstr_free(vstr);
181 }
182
183 return mp_obj_str_builder_end(o);
184
185wrong_args:
186 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "wrong number of arguments"));
187}
188
Damien George55baff42014-01-21 21:40:13 +0000189// like strstr but with specified length and allows \0 bytes
190// TODO replace with something more efficient/standard
xbe17a5a832014-03-23 23:31:58 -0700191STATIC const byte *find_subbytes(const byte *haystack, machine_uint_t hlen, const byte *needle, machine_uint_t nlen, machine_int_t direction) {
Damien George55baff42014-01-21 21:40:13 +0000192 if (hlen >= nlen) {
xbe17a5a832014-03-23 23:31:58 -0700193 machine_uint_t str_index, str_index_end;
194 if (direction > 0) {
195 str_index = 0;
196 str_index_end = hlen - nlen;
197 } else {
198 str_index = hlen - nlen;
199 str_index_end = 0;
200 }
201 for (;;) {
202 if (memcmp(&haystack[str_index], needle, nlen) == 0) {
203 //found
204 return haystack + str_index;
Damien George55baff42014-01-21 21:40:13 +0000205 }
xbe17a5a832014-03-23 23:31:58 -0700206 if (str_index == str_index_end) {
207 //not found
208 break;
Damien George55baff42014-01-21 21:40:13 +0000209 }
xbe17a5a832014-03-23 23:31:58 -0700210 str_index += direction;
Damien George55baff42014-01-21 21:40:13 +0000211 }
212 }
213 return NULL;
214}
215
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200216STATIC 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 +0000217 GET_STR_DATA_LEN(lhs_in, lhs_data, lhs_len);
Damiend99b0522013-12-21 18:17:45 +0000218 switch (op) {
Damien Georged17926d2014-03-30 13:35:08 +0100219 case MP_BINARY_OP_SUBSCR:
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200220 // TODO: need predicate to check for int-like type (bools are such for example)
221 // ["no", "yes"][1 == 2] is common idiom
222 if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
xbe9e1e8cd2014-03-12 22:57:16 -0700223 uint index = mp_get_index(mp_obj_get_type(lhs_in), lhs_len, rhs_in, false);
Damien George3e1a5c12014-03-29 13:43:38 +0000224 if (MP_OBJ_IS_TYPE(lhs_in, &mp_type_bytes)) {
Damien George7c9c6672014-01-25 00:17:36 +0000225 return MP_OBJ_NEW_SMALL_INT((mp_small_int_t)lhs_data[index]);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200226 } else {
227 return mp_obj_new_str(lhs_data + index, 1, true);
228 }
Paul Sokolovskye606cb62014-01-04 01:34:23 +0200229#if MICROPY_ENABLE_SLICE
Damien George3e1a5c12014-03-29 13:43:38 +0000230 } else if (MP_OBJ_IS_TYPE(rhs_in, &mp_type_slice)) {
Paul Sokolovsky7364af22014-02-02 02:38:22 +0200231 machine_uint_t start, stop;
Paul Sokolovskyea2509d2014-02-02 08:57:05 +0200232 if (!m_seq_get_fast_slice_indexes(lhs_len, rhs_in, &start, &stop)) {
233 assert(0);
234 }
Damien George5fa93b62014-01-22 14:35:10 +0000235 return mp_obj_new_str(lhs_data + start, stop - start, false);
Paul Sokolovskye606cb62014-01-04 01:34:23 +0200236#endif
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200237 } else {
Paul Sokolovskyf8b9d3c2014-01-04 01:38:26 +0200238 // Message doesn't match CPython, but we don't have so much bytes as they
239 // to spend them on verbose wording
Damien Georgec5966122014-02-15 16:10:44 +0000240 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "index must be int"));
Paul Sokolovsky31ba60f2014-01-03 02:51:16 +0200241 }
Damiend99b0522013-12-21 18:17:45 +0000242
Damien Georged17926d2014-03-30 13:35:08 +0100243 case MP_BINARY_OP_ADD:
244 case MP_BINARY_OP_INPLACE_ADD:
Damien George5fa93b62014-01-22 14:35:10 +0000245 if (MP_OBJ_IS_STR(rhs_in)) {
Damiend99b0522013-12-21 18:17:45 +0000246 // add 2 strings
Damien George5fa93b62014-01-22 14:35:10 +0000247
248 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
Damien George55baff42014-01-21 21:40:13 +0000249 int alloc_len = lhs_len + rhs_len;
Damien George5fa93b62014-01-22 14:35:10 +0000250
251 /* code for making qstr
Damien George55baff42014-01-21 21:40:13 +0000252 byte *q_ptr;
253 byte *val = qstr_build_start(alloc_len, &q_ptr);
254 memcpy(val, lhs_data, lhs_len);
255 memcpy(val + lhs_len, rhs_data, rhs_len);
Damien George5fa93b62014-01-22 14:35:10 +0000256 return MP_OBJ_NEW_QSTR(qstr_build_end(q_ptr));
257 */
258
259 // code for non-qstr
260 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200261 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 +0000262 memcpy(data, lhs_data, lhs_len);
263 memcpy(data + lhs_len, rhs_data, rhs_len);
264 return mp_obj_str_builder_end(s);
Damiend99b0522013-12-21 18:17:45 +0000265 }
266 break;
Damien George5fa93b62014-01-22 14:35:10 +0000267
Damien Georged17926d2014-03-30 13:35:08 +0100268 case MP_BINARY_OP_IN:
John R. Lentonc1bef212014-01-11 12:39:33 +0000269 /* NOTE `a in b` is `b.__contains__(a)` */
Damien George5fa93b62014-01-22 14:35:10 +0000270 if (MP_OBJ_IS_STR(rhs_in)) {
271 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
xbe17a5a832014-03-23 23:31:58 -0700272 return MP_BOOL(find_subbytes(lhs_data, lhs_len, rhs_data, rhs_len, 1) != NULL);
John R. Lentonc1bef212014-01-11 12:39:33 +0000273 }
274 break;
Damien George5fa93b62014-01-22 14:35:10 +0000275
Damien Georged17926d2014-03-30 13:35:08 +0100276 case MP_BINARY_OP_MULTIPLY:
Paul Sokolovsky545591a2014-01-21 00:27:33 +0200277 {
278 if (!MP_OBJ_IS_SMALL_INT(rhs_in)) {
279 return NULL;
280 }
281 int n = MP_OBJ_SMALL_INT_VALUE(rhs_in);
Damien George5fa93b62014-01-22 14:35:10 +0000282 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200283 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 +0000284 mp_seq_multiply(lhs_data, sizeof(*lhs_data), lhs_len, n, data);
285 return mp_obj_str_builder_end(s);
Paul Sokolovsky545591a2014-01-21 00:27:33 +0200286 }
Paul Sokolovsky87e85b72014-02-02 08:24:07 +0200287
Paul Sokolovsky4db727a2014-03-31 21:18:28 +0300288 case MP_BINARY_OP_MODULO: {
289 mp_obj_t *args;
290 uint n_args;
291 if (MP_OBJ_IS_TYPE(rhs_in, &mp_type_tuple)) {
292 // TODO: Support tuple subclasses?
293 mp_obj_tuple_get(rhs_in, &n_args, &args);
294 } else {
295 args = &rhs_in;
296 n_args = 1;
297 }
298 return str_modulo_format(lhs_in, n_args, args);
299 }
300
Damien Georged17926d2014-03-30 13:35:08 +0100301 // These 2 are never passed here, dealt with as a special case in mp_binary_op().
302 //case MP_BINARY_OP_EQUAL:
303 //case MP_BINARY_OP_NOT_EQUAL:
304 case MP_BINARY_OP_LESS:
305 case MP_BINARY_OP_LESS_EQUAL:
306 case MP_BINARY_OP_MORE:
307 case MP_BINARY_OP_MORE_EQUAL:
Paul Sokolovsky87e85b72014-02-02 08:24:07 +0200308 if (MP_OBJ_IS_STR(rhs_in)) {
309 GET_STR_DATA_LEN(rhs_in, rhs_data, rhs_len);
310 return MP_BOOL(mp_seq_cmp_bytes(op, lhs_data, lhs_len, rhs_data, rhs_len));
311 }
Damiend99b0522013-12-21 18:17:45 +0000312 }
313
314 return MP_OBJ_NULL; // op not supported
315}
316
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200317STATIC mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
Damien George5fa93b62014-01-22 14:35:10 +0000318 assert(MP_OBJ_IS_STR(self_in));
Damiend99b0522013-12-21 18:17:45 +0000319
Damien Georgefe8fb912014-01-02 16:36:09 +0000320 // get separation string
Damien George5fa93b62014-01-22 14:35:10 +0000321 GET_STR_DATA_LEN(self_in, sep_str, sep_len);
Damien Georgefe8fb912014-01-02 16:36:09 +0000322
323 // process args
Damiend99b0522013-12-21 18:17:45 +0000324 uint seq_len;
325 mp_obj_t *seq_items;
Damien George07ddab52014-03-29 13:15:08 +0000326 if (MP_OBJ_IS_TYPE(arg, &mp_type_tuple)) {
Damiend99b0522013-12-21 18:17:45 +0000327 mp_obj_tuple_get(arg, &seq_len, &seq_items);
Damien George3e1a5c12014-03-29 13:43:38 +0000328 } else if (MP_OBJ_IS_TYPE(arg, &mp_type_list)) {
Damiend99b0522013-12-21 18:17:45 +0000329 mp_obj_list_get(arg, &seq_len, &seq_items);
330 } else {
331 goto bad_arg;
332 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000333
334 // count required length
335 int required_len = 0;
Damiend99b0522013-12-21 18:17:45 +0000336 for (int i = 0; i < seq_len; i++) {
Damien George5fa93b62014-01-22 14:35:10 +0000337 if (!MP_OBJ_IS_STR(seq_items[i])) {
Damiend99b0522013-12-21 18:17:45 +0000338 goto bad_arg;
339 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000340 if (i > 0) {
341 required_len += sep_len;
342 }
Damien George5fa93b62014-01-22 14:35:10 +0000343 GET_STR_LEN(seq_items[i], l);
344 required_len += l;
Damiend99b0522013-12-21 18:17:45 +0000345 }
346
347 // make joined string
Damien George5fa93b62014-01-22 14:35:10 +0000348 byte *data;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +0200349 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 +0000350 for (int i = 0; i < seq_len; i++) {
Damiend99b0522013-12-21 18:17:45 +0000351 if (i > 0) {
Damien George5fa93b62014-01-22 14:35:10 +0000352 memcpy(data, sep_str, sep_len);
353 data += sep_len;
Damiend99b0522013-12-21 18:17:45 +0000354 }
Damien George5fa93b62014-01-22 14:35:10 +0000355 GET_STR_DATA_LEN(seq_items[i], s, l);
356 memcpy(data, s, l);
357 data += l;
Damiend99b0522013-12-21 18:17:45 +0000358 }
Damien Georgefe8fb912014-01-02 16:36:09 +0000359
360 // return joined string
Damien George5fa93b62014-01-22 14:35:10 +0000361 return mp_obj_str_builder_end(joined_str);
Damiend99b0522013-12-21 18:17:45 +0000362
363bad_arg:
Damien Georgec5966122014-02-15 16:10:44 +0000364 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "?str.join expecting a list of str's"));
Damiend99b0522013-12-21 18:17:45 +0000365}
366
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200367#define is_ws(c) ((c) == ' ' || (c) == '\t')
368
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200369STATIC mp_obj_t str_split(uint n_args, const mp_obj_t *args) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200370 int splits = -1;
371 mp_obj_t sep = mp_const_none;
372 if (n_args > 1) {
373 sep = args[1];
374 if (n_args > 2) {
375 splits = MP_OBJ_SMALL_INT_VALUE(args[2]);
376 }
377 }
378 assert(sep == mp_const_none);
Damien George12eacca2014-01-21 21:54:15 +0000379 (void)sep; // unused; to hush compiler warning
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200380 mp_obj_t res = mp_obj_new_list(0, NULL);
Damien George5fa93b62014-01-22 14:35:10 +0000381 GET_STR_DATA_LEN(args[0], s, len);
382 const byte *top = s + len;
383 const byte *start;
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200384
385 // Initial whitespace is not counted as split, so we pre-do it
Damien George5fa93b62014-01-22 14:35:10 +0000386 while (s < top && is_ws(*s)) s++;
387 while (s < top && splits != 0) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200388 start = s;
Damien George5fa93b62014-01-22 14:35:10 +0000389 while (s < top && !is_ws(*s)) s++;
Damien George15d18062014-03-31 16:28:13 +0100390 mp_obj_list_append(res, mp_obj_new_str(start, s - start, false));
Damien George5fa93b62014-01-22 14:35:10 +0000391 if (s >= top) {
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200392 break;
393 }
Damien George5fa93b62014-01-22 14:35:10 +0000394 while (s < top && is_ws(*s)) s++;
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200395 if (splits > 0) {
396 splits--;
397 }
398 }
399
Damien George5fa93b62014-01-22 14:35:10 +0000400 if (s < top) {
Damien George15d18062014-03-31 16:28:13 +0100401 mp_obj_list_append(res, mp_obj_new_str(s, top - s, false));
Paul Sokolovsky4c316552014-01-21 05:00:21 +0200402 }
403
404 return res;
405}
406
xbe17a5a832014-03-23 23:31:58 -0700407STATIC mp_obj_t str_finder(uint n_args, const mp_obj_t *args, machine_int_t direction) {
John R. Lentone8204912014-01-12 21:53:52 +0000408 assert(2 <= n_args && n_args <= 4);
Damien George5fa93b62014-01-22 14:35:10 +0000409 assert(MP_OBJ_IS_STR(args[0]));
410 assert(MP_OBJ_IS_STR(args[1]));
John R. Lentone8204912014-01-12 21:53:52 +0000411
Damien George5fa93b62014-01-22 14:35:10 +0000412 GET_STR_DATA_LEN(args[0], haystack, haystack_len);
413 GET_STR_DATA_LEN(args[1], needle, needle_len);
John R. Lentone8204912014-01-12 21:53:52 +0000414
xbec5538882014-03-16 17:58:35 -0700415 machine_uint_t start = 0;
416 machine_uint_t end = haystack_len;
John R. Lentone8204912014-01-12 21:53:52 +0000417 if (n_args >= 3 && args[2] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +0000418 start = mp_get_index(&mp_type_str, haystack_len, args[2], true);
John R. Lentone8204912014-01-12 21:53:52 +0000419 }
420 if (n_args >= 4 && args[3] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +0000421 end = mp_get_index(&mp_type_str, haystack_len, args[3], true);
John R. Lentone8204912014-01-12 21:53:52 +0000422 }
423
xbe17a5a832014-03-23 23:31:58 -0700424 const byte *p = find_subbytes(haystack + start, end - start, needle, needle_len, direction);
Damien George23005372014-01-13 19:39:01 +0000425 if (p == NULL) {
426 // not found
427 return MP_OBJ_NEW_SMALL_INT(-1);
428 } else {
429 // found
xbe17a5a832014-03-23 23:31:58 -0700430 return MP_OBJ_NEW_SMALL_INT(p - haystack);
John R. Lentone8204912014-01-12 21:53:52 +0000431 }
John R. Lentone8204912014-01-12 21:53:52 +0000432}
433
xbe17a5a832014-03-23 23:31:58 -0700434STATIC mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
435 return str_finder(n_args, args, 1);
436}
437
438STATIC mp_obj_t str_rfind(uint n_args, const mp_obj_t *args) {
439 return str_finder(n_args, args, -1);
440}
441
Paul Sokolovsky1eacefe2014-01-23 01:20:40 +0200442// TODO: (Much) more variety in args
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200443STATIC mp_obj_t str_startswith(mp_obj_t self_in, mp_obj_t arg) {
Paul Sokolovsky1eacefe2014-01-23 01:20:40 +0200444 GET_STR_DATA_LEN(self_in, str, str_len);
445 GET_STR_DATA_LEN(arg, prefix, prefix_len);
446 if (prefix_len > str_len) {
447 return mp_const_false;
448 }
449 return MP_BOOL(memcmp(str, prefix, prefix_len) == 0);
450}
451
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200452STATIC mp_obj_t str_strip(uint n_args, const mp_obj_t *args) {
xbe7b0f39f2014-01-08 14:23:45 -0800453 assert(1 <= n_args && n_args <= 2);
Damien George5fa93b62014-01-22 14:35:10 +0000454 assert(MP_OBJ_IS_STR(args[0]));
455
456 const byte *chars_to_del;
457 uint chars_to_del_len;
458 static const byte whitespace[] = " \t\n\r\v\f";
xbe7b0f39f2014-01-08 14:23:45 -0800459
460 if (n_args == 1) {
461 chars_to_del = whitespace;
Damien George5fa93b62014-01-22 14:35:10 +0000462 chars_to_del_len = sizeof(whitespace);
xbe7b0f39f2014-01-08 14:23:45 -0800463 } else {
Damien George5fa93b62014-01-22 14:35:10 +0000464 assert(MP_OBJ_IS_STR(args[1]));
465 GET_STR_DATA_LEN(args[1], s, l);
466 chars_to_del = s;
467 chars_to_del_len = l;
xbe7b0f39f2014-01-08 14:23:45 -0800468 }
469
Damien George5fa93b62014-01-22 14:35:10 +0000470 GET_STR_DATA_LEN(args[0], orig_str, orig_str_len);
xbe7b0f39f2014-01-08 14:23:45 -0800471
xbec5538882014-03-16 17:58:35 -0700472 machine_uint_t first_good_char_pos = 0;
xbe7b0f39f2014-01-08 14:23:45 -0800473 bool first_good_char_pos_set = false;
xbec5538882014-03-16 17:58:35 -0700474 machine_uint_t last_good_char_pos = 0;
475 for (machine_uint_t i = 0; i < orig_str_len; i++) {
xbe17a5a832014-03-23 23:31:58 -0700476 if (find_subbytes(chars_to_del, chars_to_del_len, &orig_str[i], 1, 1) == NULL) {
xbe7b0f39f2014-01-08 14:23:45 -0800477 last_good_char_pos = i;
478 if (!first_good_char_pos_set) {
479 first_good_char_pos = i;
480 first_good_char_pos_set = true;
481 }
482 }
483 }
484
485 if (first_good_char_pos == 0 && last_good_char_pos == 0) {
Damien George5fa93b62014-01-22 14:35:10 +0000486 // string is all whitespace, return ''
487 return MP_OBJ_NEW_QSTR(MP_QSTR_);
xbe7b0f39f2014-01-08 14:23:45 -0800488 }
489
490 assert(last_good_char_pos >= first_good_char_pos);
491 //+1 to accomodate the last character
xbec5538882014-03-16 17:58:35 -0700492 machine_uint_t stripped_len = last_good_char_pos - first_good_char_pos + 1;
Damien George5fa93b62014-01-22 14:35:10 +0000493 return mp_obj_new_str(orig_str + first_good_char_pos, stripped_len, false);
xbe7b0f39f2014-01-08 14:23:45 -0800494}
495
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700496// Takes an int arg, but only parses unsigned numbers, and only changes
497// *num if at least one digit was parsed.
498static int str_to_int(const char *str, int *num) {
499 const char *s = str;
500 if (unichar_isdigit(*s)) {
501 *num = 0;
502 do {
503 *num = *num * 10 + (*s - '0');
504 s++;
505 }
506 while (unichar_isdigit(*s));
507 }
508 return s - str;
509}
510
511static bool isalignment(char ch) {
512 return ch && strchr("<>=^", ch) != NULL;
513}
514
515static bool istype(char ch) {
516 return ch && strchr("bcdeEfFgGnosxX%", ch) != NULL;
517}
518
519static bool arg_looks_integer(mp_obj_t arg) {
520 return MP_OBJ_IS_TYPE(arg, &mp_type_bool) || MP_OBJ_IS_INT(arg);
521}
522
523static bool arg_looks_numeric(mp_obj_t arg) {
524 return arg_looks_integer(arg)
525#if MICROPY_ENABLE_FLOAT
526 || MP_OBJ_IS_TYPE(arg, &mp_type_float)
527#endif
528 ;
529}
530
Damien Georgea11ceca2014-01-19 16:02:09 +0000531mp_obj_t str_format(uint n_args, const mp_obj_t *args) {
Damien George5fa93b62014-01-22 14:35:10 +0000532 assert(MP_OBJ_IS_STR(args[0]));
Damiend99b0522013-12-21 18:17:45 +0000533
Damien George5fa93b62014-01-22 14:35:10 +0000534 GET_STR_DATA_LEN(args[0], str, len);
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700535 int arg_i = 0;
Damiend99b0522013-12-21 18:17:45 +0000536 vstr_t *vstr = vstr_new();
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700537 pfenv_t pfenv_vstr;
538 pfenv_vstr.data = vstr;
539 pfenv_vstr.print_strn = pfenv_vstr_add_strn;
540
Damien George5fa93b62014-01-22 14:35:10 +0000541 for (const byte *top = str + len; str < top; str++) {
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700542 if (*str == '}') {
Damiend99b0522013-12-21 18:17:45 +0000543 str++;
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700544 if (str < top && *str == '}') {
545 vstr_add_char(vstr, '}');
546 continue;
547 }
548 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "Single '}' encountered in format string"));
549 }
550 if (*str != '{') {
551 vstr_add_char(vstr, *str);
552 continue;
553 }
554
555 str++;
556 if (str < top && *str == '{') {
557 vstr_add_char(vstr, '{');
558 continue;
559 }
560
561 // replacement_field ::= "{" [field_name] ["!" conversion] [":" format_spec] "}"
562
563 vstr_t *field_name = NULL;
564 char conversion = '\0';
565 vstr_t *format_spec = NULL;
566
567 if (str < top && *str != '}' && *str != '!' && *str != ':') {
568 field_name = vstr_new();
569 while (str < top && *str != '}' && *str != '!' && *str != ':') {
570 vstr_add_char(field_name, *str++);
571 }
572 vstr_add_char(field_name, '\0');
573 }
574
575 // conversion ::= "r" | "s"
576
577 if (str < top && *str == '!') {
578 str++;
579 if (str < top && (*str == 'r' || *str == 's')) {
580 conversion = *str++;
Paul Sokolovskyf2b796e2014-01-15 22:45:20 +0200581 } else {
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700582 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "end of format while looking for conversion specifier"));
583 }
584 }
585
586 if (str < top && *str == ':') {
587 str++;
588 // {:} is the same as {}, which is the same as {!s}
589 // This makes a difference when passing in a True or False
590 // '{}'.format(True) returns 'True'
591 // '{:d}'.format(True) returns '1'
592 // So we treat {:} as {} and this later gets treated to be {!s}
593 if (*str != '}') {
594 format_spec = vstr_new();
595 while (str < top && *str != '}') {
596 vstr_add_char(format_spec, *str++);
Damiend99b0522013-12-21 18:17:45 +0000597 }
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700598 vstr_add_char(format_spec, '\0');
599 }
600 }
601 if (str >= top) {
602 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "unmatched '{' in format"));
603 }
604 if (*str != '}') {
605 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "expected ':' after format specifier"));
606 }
607
608 mp_obj_t arg = mp_const_none;
609
610 if (field_name) {
611 if (arg_i > 0) {
612 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "cannot switch from automatic field numbering to manual field specification"));
613 }
614 int index;
615 if (str_to_int(vstr_str(field_name), &index) != vstr_len(field_name) - 1) {
616 nlr_jump(mp_obj_new_exception_msg(&mp_type_KeyError, "attributes not supported yet"));
617 }
618 if (index >= n_args - 1) {
619 nlr_jump(mp_obj_new_exception_msg(&mp_type_IndexError, "tuple index out of range"));
620 }
621 arg = args[index + 1];
622 arg_i = -1;
623 vstr_free(field_name);
624 field_name = NULL;
625 } else {
626 if (arg_i < 0) {
627 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "cannot switch from manual field specification to automatic field numbering"));
628 }
629 if (arg_i >= n_args - 1) {
630 nlr_jump(mp_obj_new_exception_msg(&mp_type_IndexError, "tuple index out of range"));
631 }
632 arg = args[arg_i + 1];
633 arg_i++;
634 }
635 if (!format_spec && !conversion) {
636 conversion = 's';
637 }
638 if (conversion) {
639 mp_print_kind_t print_kind;
640 if (conversion == 's') {
641 print_kind = PRINT_STR;
642 } else if (conversion == 'r') {
643 print_kind = PRINT_REPR;
644 } else {
645 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError, "Unknown conversion specifier %c", conversion));
646 }
647 vstr_t *arg_vstr = vstr_new();
648 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, arg_vstr, arg, print_kind);
649 arg = mp_obj_new_str((const byte *)vstr_str(arg_vstr), vstr_len(arg_vstr), false);
650 vstr_free(arg_vstr);
651 }
652
653 char sign = '\0';
654 char fill = '\0';
655 char align = '\0';
656 int width = -1;
657 int precision = -1;
658 char type = '\0';
659 int flags = 0;
660
661 if (format_spec) {
662 // The format specifier (from http://docs.python.org/2/library/string.html#formatspec)
663 //
664 // [[fill]align][sign][#][0][width][,][.precision][type]
665 // fill ::= <any character>
666 // align ::= "<" | ">" | "=" | "^"
667 // sign ::= "+" | "-" | " "
668 // width ::= integer
669 // precision ::= integer
670 // type ::= "b" | "c" | "d" | "e" | "E" | "f" | "F" | "g" | "G" | "n" | "o" | "s" | "x" | "X" | "%"
671
672 const char *s = vstr_str(format_spec);
673 if (isalignment(*s)) {
674 align = *s++;
675 } else if (*s && isalignment(s[1])) {
676 fill = *s++;
677 align = *s++;
678 }
679 if (*s == '+' || *s == '-' || *s == ' ') {
680 if (*s == '+') {
681 flags |= PF_FLAG_SHOW_SIGN;
682 } else if (*s == ' ') {
683 flags |= PF_FLAG_SPACE_SIGN;
684 }
685 sign = *s++;
686 }
687 if (*s == '#') {
688 flags |= PF_FLAG_SHOW_PREFIX;
689 s++;
690 }
691 if (*s == '0') {
692 if (!align) {
693 align = '=';
694 }
695 if (!fill) {
696 fill = '0';
697 }
698 }
699 s += str_to_int(s, &width);
700 if (*s == ',') {
701 flags |= PF_FLAG_SHOW_COMMA;
702 s++;
703 }
704 if (*s == '.') {
705 s++;
706 s += str_to_int(s, &precision);
707 }
708 if (istype(*s)) {
709 type = *s++;
710 }
711 if (*s) {
712 nlr_jump(mp_obj_new_exception_msg(&mp_type_KeyError, "Invalid conversion specification"));
713 }
714 vstr_free(format_spec);
715 format_spec = NULL;
716 }
717 if (!align) {
718 if (arg_looks_numeric(arg)) {
719 align = '>';
720 } else {
721 align = '<';
722 }
723 }
724 if (!fill) {
725 fill = ' ';
726 }
727
728 if (sign) {
729 if (type == 's') {
730 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "Sign not allowed in string format specifier"));
731 }
732 if (type == 'c') {
733 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "Sign not allowed with integer format specifier 'c'"));
Damiend99b0522013-12-21 18:17:45 +0000734 }
735 } else {
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700736 sign = '-';
737 }
738
739 switch (align) {
740 case '<': flags |= PF_FLAG_LEFT_ADJUST; break;
741 case '=': flags |= PF_FLAG_PAD_AFTER_SIGN; break;
742 case '^': flags |= PF_FLAG_CENTER_ADJUST; break;
743 }
744
745 if (arg_looks_integer(arg)) {
746 switch (type) {
747 case 'b':
748 pfenv_print_int(&pfenv_vstr, mp_obj_get_int(arg), 1, 2, 'a', flags, fill, width);
749 continue;
750
751 case 'c':
752 {
753 char ch = mp_obj_get_int(arg);
754 pfenv_print_strn(&pfenv_vstr, &ch, 1, flags, fill, width);
755 continue;
756 }
757
758 case '\0': // No explicit format type implies 'd'
759 case 'n': // I don't think we support locales in uPy so use 'd'
760 case 'd':
761 pfenv_print_int(&pfenv_vstr, mp_obj_get_int(arg), 1, 10, 'a', flags, fill, width);
762 continue;
763
764 case 'o':
765 pfenv_print_int(&pfenv_vstr, mp_obj_get_int(arg), 1, 8, 'a', flags, fill, width);
766 continue;
767
768 case 'x':
769 pfenv_print_int(&pfenv_vstr, mp_obj_get_int(arg), 1, 16, 'a', flags, fill, width);
770 continue;
771
772 case 'X':
773 pfenv_print_int(&pfenv_vstr, mp_obj_get_int(arg), 1, 16, 'A', flags, fill, width);
774 continue;
775
776 case 'e':
777 case 'E':
778 case 'f':
779 case 'F':
780 case 'g':
781 case 'G':
782 case '%':
783 // The floating point formatters all work with anything that
784 // looks like an integer
785 break;
786
787 default:
788 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError,
789 "Unknown format code '%c' for object of type '%s'", type, mp_obj_get_type_str(arg)));
790 }
791 }
792#if MICROPY_ENABLE_FLOAT
793 if (arg_looks_numeric(arg)) {
794 if (!type) {
795
796 // Even though the docs say that an unspecified type is the same
797 // as 'g', there is one subtle difference, when the exponent
798 // is one less than the precision.
799 //
800 // '{:10.1}'.format(0.0) ==> '0e+00'
801 // '{:10.1g}'.format(0.0) ==> '0'
802 //
803 // TODO: Figure out how to deal with this.
804 //
805 // A proper solution would involve adding a special flag
806 // or something to format_float, and create a format_double
807 // to deal with doubles. In order to fix this when using
808 // sprintf, we'd need to use the e format and tweak the
809 // returned result to strip trailing zeros like the g format
810 // does.
811 //
812 // {:10.3} and {:10.2e} with 1.23e2 both produce 1.23e+02
813 // but with 1.e2 you get 1e+02 and 1.00e+02
814 //
815 // Stripping the trailing 0's (like g) does would make the
816 // e format give us the right format.
817 //
818 // CPython sources say:
819 // Omitted type specifier. Behaves in the same way as repr(x)
820 // and str(x) if no precision is given, else like 'g', but with
821 // at least one digit after the decimal point. */
822
823 type = 'g';
824 }
825 if (type == 'n') {
826 type = 'g';
827 }
828
829 flags |= PF_FLAG_PAD_NAN_INF; // '{:06e}'.format(float('-inf')) should give '-00inf'
830 switch (type) {
831 case 'e':
832 case 'E':
833 case 'f':
834 case 'F':
835 case 'g':
836 case 'G':
837 pfenv_print_float(&pfenv_vstr, mp_obj_get_float(arg), type, flags, fill, width, precision);
838 break;
839
840 case '%':
841 flags |= PF_FLAG_ADD_PERCENT;
842 pfenv_print_float(&pfenv_vstr, mp_obj_get_float(arg) * 100.0F, 'f', flags, fill, width, precision);
843 break;
844
845 default:
846 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError,
847 "Unknown format code '%c' for object of type 'float'",
848 type, mp_obj_get_type_str(arg)));
849 }
850#endif
851 } else {
852 if (align == '=') {
853 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "'=' alignment not allowed in string format specifier"));
854 }
855 switch (type) {
856 case '\0':
857 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, arg, PRINT_STR);
858 break;
859
860 case 's':
861 {
862 uint len;
863 const char *s = mp_obj_str_get_data(arg, &len);
864 if (precision < 0) {
865 precision = len;
866 }
867 if (len > precision) {
868 len = precision;
869 }
870 pfenv_print_strn(&pfenv_vstr, s, len, flags, fill, width);
871 break;
872 }
873
874 default:
875 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError,
876 "Unknown format code '%c' for object of type 'str'",
877 type, mp_obj_get_type_str(arg)));
878 }
Damiend99b0522013-12-21 18:17:45 +0000879 }
880 }
881
Damien George5fa93b62014-01-22 14:35:10 +0000882 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
883 vstr_free(vstr);
884 return s;
Damiend99b0522013-12-21 18:17:45 +0000885}
886
Paul Sokolovsky4db727a2014-03-31 21:18:28 +0300887STATIC mp_obj_t str_modulo_format(mp_obj_t pattern, uint n_args, const mp_obj_t *args) {
888 assert(MP_OBJ_IS_STR(pattern));
889
890 GET_STR_DATA_LEN(pattern, str, len);
891 int arg_i = 0;
892 vstr_t *vstr = vstr_new();
893 for (const byte *top = str + len; str < top; str++) {
894 if (*str == '%') {
895 if (++str >= top) {
896 break;
897 }
898 if (*str == '%') {
899 vstr_add_char(vstr, '%');
900 } else {
901 if (arg_i >= n_args) {
902 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "not enough arguments for format string"));
903 }
904 switch (*str) {
905 case 's':
906 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR);
907 break;
908 case 'r':
909 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_REPR);
910 break;
911 }
912 arg_i++;
913 }
914 } else {
915 vstr_add_char(vstr, *str);
916 }
917 }
918
919 if (arg_i != n_args) {
920 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "not all arguments converted during string formatting"));
921 }
922
923 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
924 vstr_free(vstr);
925 return s;
926}
927
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200928STATIC mp_obj_t str_replace(uint n_args, const mp_obj_t *args) {
xbe480c15a2014-01-30 22:17:30 -0800929 assert(MP_OBJ_IS_STR(args[0]));
930 assert(MP_OBJ_IS_STR(args[1]));
931 assert(MP_OBJ_IS_STR(args[2]));
932
Damien George94f68302014-01-31 23:45:12 +0000933 machine_int_t max_rep = 0;
xbe480c15a2014-01-30 22:17:30 -0800934 if (n_args == 4) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200935 assert(MP_OBJ_IS_SMALL_INT(args[3]));
936 max_rep = MP_OBJ_SMALL_INT_VALUE(args[3]);
937 if (max_rep == 0) {
938 return args[0];
939 } else if (max_rep < 0) {
940 max_rep = 0;
941 }
xbe480c15a2014-01-30 22:17:30 -0800942 }
Damien George94f68302014-01-31 23:45:12 +0000943
944 // if max_rep is still 0 by this point we will need to do all possible replacements
xbe480c15a2014-01-30 22:17:30 -0800945
946 GET_STR_DATA_LEN(args[0], str, str_len);
947 GET_STR_DATA_LEN(args[1], old, old_len);
948 GET_STR_DATA_LEN(args[2], new, new_len);
Damien George94f68302014-01-31 23:45:12 +0000949
950 // old won't exist in str if it's longer, so nothing to replace
xbe480c15a2014-01-30 22:17:30 -0800951 if (old_len > str_len) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200952 return args[0];
xbe480c15a2014-01-30 22:17:30 -0800953 }
954
Damien George94f68302014-01-31 23:45:12 +0000955 // data for the replaced string
956 byte *data = NULL;
957 mp_obj_t replaced_str = MP_OBJ_NULL;
xbe480c15a2014-01-30 22:17:30 -0800958
Damien George94f68302014-01-31 23:45:12 +0000959 // do 2 passes over the string:
960 // first pass computes the required length of the replaced string
961 // second pass does the replacements
962 for (;;) {
963 machine_uint_t replaced_str_index = 0;
964 machine_uint_t num_replacements_done = 0;
965 const byte *old_occurrence;
966 const byte *offset_ptr = str;
967 machine_uint_t offset_num = 0;
xbe17a5a832014-03-23 23:31:58 -0700968 while ((old_occurrence = find_subbytes(offset_ptr, str_len - offset_num, old, old_len, 1)) != NULL) {
Damien George94f68302014-01-31 23:45:12 +0000969 // copy from just after end of last occurrence of to-be-replaced string to right before start of next occurrence
970 if (data != NULL) {
971 memcpy(data + replaced_str_index, offset_ptr, old_occurrence - offset_ptr);
972 }
973 replaced_str_index += old_occurrence - offset_ptr;
974 // copy the replacement string
975 if (data != NULL) {
976 memcpy(data + replaced_str_index, new, new_len);
977 }
978 replaced_str_index += new_len;
979 offset_ptr = old_occurrence + old_len;
980 offset_num = offset_ptr - str;
xbe480c15a2014-01-30 22:17:30 -0800981
Damien George94f68302014-01-31 23:45:12 +0000982 num_replacements_done++;
983 if (max_rep != 0 && num_replacements_done == max_rep){
984 break;
985 }
986 }
987
988 // copy from just after end of last occurrence of to-be-replaced string to end of old string
989 if (data != NULL) {
990 memcpy(data + replaced_str_index, offset_ptr, str_len - offset_num);
991 }
992 replaced_str_index += str_len - offset_num;
993
994 if (data == NULL) {
995 // first pass
996 if (num_replacements_done == 0) {
997 // no substr found, return original string
998 return args[0];
999 } else {
1000 // substr found, allocate new string
1001 replaced_str = mp_obj_str_builder_start(mp_obj_get_type(args[0]), replaced_str_index, &data);
1002 }
1003 } else {
1004 // second pass, we are done
1005 break;
1006 }
xbe480c15a2014-01-30 22:17:30 -08001007 }
Damien George94f68302014-01-31 23:45:12 +00001008
xbe480c15a2014-01-30 22:17:30 -08001009 return mp_obj_str_builder_end(replaced_str);
1010}
1011
xbe9e1e8cd2014-03-12 22:57:16 -07001012STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
1013 assert(2 <= n_args && n_args <= 4);
1014 assert(MP_OBJ_IS_STR(args[0]));
1015 assert(MP_OBJ_IS_STR(args[1]));
1016
1017 GET_STR_DATA_LEN(args[0], haystack, haystack_len);
1018 GET_STR_DATA_LEN(args[1], needle, needle_len);
1019
Damien George536dde22014-03-13 22:07:55 +00001020 machine_uint_t start = 0;
1021 machine_uint_t end = haystack_len;
xbe9e1e8cd2014-03-12 22:57:16 -07001022 if (n_args >= 3 && args[2] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +00001023 start = mp_get_index(&mp_type_str, haystack_len, args[2], true);
xbe9e1e8cd2014-03-12 22:57:16 -07001024 }
1025 if (n_args >= 4 && args[3] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +00001026 end = mp_get_index(&mp_type_str, haystack_len, args[3], true);
xbe9e1e8cd2014-03-12 22:57:16 -07001027 }
1028
Damien George536dde22014-03-13 22:07:55 +00001029 // if needle_len is zero then we count each gap between characters as an occurrence
1030 if (needle_len == 0) {
1031 return MP_OBJ_NEW_SMALL_INT(end - start + 1);
xbe9e1e8cd2014-03-12 22:57:16 -07001032 }
1033
Damien George536dde22014-03-13 22:07:55 +00001034 // count the occurrences
1035 machine_int_t num_occurrences = 0;
xbec5d70ba2014-03-13 00:29:15 -07001036 for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
1037 if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
1038 num_occurrences++;
1039 haystack_index += needle_len - 1;
1040 }
xbe9e1e8cd2014-03-12 22:57:16 -07001041 }
1042
1043 return MP_OBJ_NEW_SMALL_INT(num_occurrences);
1044}
1045
Damien Georgeb035db32014-03-21 20:39:40 +00001046STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, machine_int_t direction) {
xbe613a8e32014-03-18 00:06:29 -07001047 assert(MP_OBJ_IS_STR(self_in));
1048 if (!MP_OBJ_IS_STR(arg)) {
1049 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
1050 "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
1051 }
Damien Georgeb035db32014-03-21 20:39:40 +00001052
xbe613a8e32014-03-18 00:06:29 -07001053 GET_STR_DATA_LEN(self_in, str, str_len);
1054 GET_STR_DATA_LEN(arg, sep, sep_len);
1055
1056 if (sep_len == 0) {
1057 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
1058 }
Damien Georgeb035db32014-03-21 20:39:40 +00001059
1060 mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
1061
1062 if (direction > 0) {
1063 result[0] = self_in;
xbe0a6894c2014-03-21 01:12:26 -07001064 } else {
Damien Georgeb035db32014-03-21 20:39:40 +00001065 result[2] = self_in;
xbe0a6894c2014-03-21 01:12:26 -07001066 }
xbe613a8e32014-03-18 00:06:29 -07001067
xbe17a5a832014-03-23 23:31:58 -07001068 const byte *position_ptr = find_subbytes(str, str_len, sep, sep_len, direction);
1069 if (position_ptr != NULL) {
1070 machine_uint_t position = position_ptr - str;
1071 result[0] = mp_obj_new_str(str, position, false);
1072 result[1] = arg;
1073 result[2] = mp_obj_new_str(str + position + sep_len, str_len - position - sep_len, false);
xbe613a8e32014-03-18 00:06:29 -07001074 }
Damien Georgeb035db32014-03-21 20:39:40 +00001075
xbe0a6894c2014-03-21 01:12:26 -07001076 return mp_obj_new_tuple(3, result);
xbe613a8e32014-03-18 00:06:29 -07001077}
1078
Damien Georgeb035db32014-03-21 20:39:40 +00001079STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
1080 return str_partitioner(self_in, arg, 1);
xbe0a6894c2014-03-21 01:12:26 -07001081}
xbe4504ea82014-03-19 00:46:14 -07001082
Damien Georgeb035db32014-03-21 20:39:40 +00001083STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg) {
1084 return str_partitioner(self_in, arg, -1);
xbe4504ea82014-03-19 00:46:14 -07001085}
1086
Damien George2da98302014-03-09 19:58:18 +00001087STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
1088 if (flags == BUFFER_READ) {
1089 GET_STR_DATA_LEN(self_in, str_data, str_len);
1090 bufinfo->buf = (void*)str_data;
1091 bufinfo->len = str_len;
1092 return 0;
1093 } else {
1094 // can't write to a string
1095 bufinfo->buf = NULL;
1096 bufinfo->len = 0;
1097 return 1;
1098 }
1099}
1100
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001101STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
xbe17a5a832014-03-23 23:31:58 -07001102STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rfind_obj, 2, 4, str_rfind);
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001103STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
1104STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
1105STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
1106STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
1107STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
1108STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
xbe9e1e8cd2014-03-12 22:57:16 -07001109STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
xbe613a8e32014-03-18 00:06:29 -07001110STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_partition_obj, str_partition);
xbe4504ea82014-03-19 00:46:14 -07001111STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition);
Damiend99b0522013-12-21 18:17:45 +00001112
Damien George9b196cd2014-03-26 21:47:19 +00001113STATIC const mp_map_elem_t str_locals_dict_table[] = {
1114 { MP_OBJ_NEW_QSTR(MP_QSTR_find), (mp_obj_t)&str_find_obj },
1115 { MP_OBJ_NEW_QSTR(MP_QSTR_rfind), (mp_obj_t)&str_rfind_obj },
1116 { MP_OBJ_NEW_QSTR(MP_QSTR_join), (mp_obj_t)&str_join_obj },
1117 { MP_OBJ_NEW_QSTR(MP_QSTR_split), (mp_obj_t)&str_split_obj },
1118 { MP_OBJ_NEW_QSTR(MP_QSTR_startswith), (mp_obj_t)&str_startswith_obj },
1119 { MP_OBJ_NEW_QSTR(MP_QSTR_strip), (mp_obj_t)&str_strip_obj },
1120 { MP_OBJ_NEW_QSTR(MP_QSTR_format), (mp_obj_t)&str_format_obj },
1121 { MP_OBJ_NEW_QSTR(MP_QSTR_replace), (mp_obj_t)&str_replace_obj },
1122 { MP_OBJ_NEW_QSTR(MP_QSTR_count), (mp_obj_t)&str_count_obj },
1123 { MP_OBJ_NEW_QSTR(MP_QSTR_partition), (mp_obj_t)&str_partition_obj },
1124 { MP_OBJ_NEW_QSTR(MP_QSTR_rpartition), (mp_obj_t)&str_rpartition_obj },
ian-v7a16fad2014-01-06 09:52:29 -08001125};
Damien George97209d32014-01-07 15:58:30 +00001126
Damien George9b196cd2014-03-26 21:47:19 +00001127STATIC MP_DEFINE_CONST_DICT(str_locals_dict, str_locals_dict_table);
1128
Damien George3e1a5c12014-03-29 13:43:38 +00001129const mp_obj_type_t mp_type_str = {
Damien Georgec5966122014-02-15 16:10:44 +00001130 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001131 .name = MP_QSTR_str,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001132 .print = str_print,
Paul Sokolovskybe020c22014-03-21 11:39:01 +02001133 .make_new = str_make_new,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001134 .binary_op = str_binary_op,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001135 .getiter = mp_obj_new_str_iterator,
Damien George2da98302014-03-09 19:58:18 +00001136 .buffer_p = { .get_buffer = str_get_buffer },
Damien George9b196cd2014-03-26 21:47:19 +00001137 .locals_dict = (mp_obj_t)&str_locals_dict,
Damiend99b0522013-12-21 18:17:45 +00001138};
1139
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001140// Reuses most of methods from str
Damien George3e1a5c12014-03-29 13:43:38 +00001141const mp_obj_type_t mp_type_bytes = {
Damien Georgec5966122014-02-15 16:10:44 +00001142 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001143 .name = MP_QSTR_bytes,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001144 .print = str_print,
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001145 .make_new = bytes_make_new,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001146 .binary_op = str_binary_op,
1147 .getiter = mp_obj_new_bytes_iterator,
Damien George9b196cd2014-03-26 21:47:19 +00001148 .locals_dict = (mp_obj_t)&str_locals_dict,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001149};
1150
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001151// the zero-length bytes
Damien George3e1a5c12014-03-29 13:43:38 +00001152STATIC const mp_obj_str_t empty_bytes_obj = {{&mp_type_bytes}, 0, 0, NULL};
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001153const mp_obj_t mp_const_empty_bytes = (mp_obj_t)&empty_bytes_obj;
1154
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001155mp_obj_t mp_obj_str_builder_start(const mp_obj_type_t *type, uint len, byte **data) {
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001156 mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001157 o->base.type = type;
Damien George5fa93b62014-01-22 14:35:10 +00001158 o->len = len;
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001159 byte *p = m_new(byte, len + 1);
1160 o->data = p;
1161 *data = p;
Damiend99b0522013-12-21 18:17:45 +00001162 return o;
1163}
1164
Damien George5fa93b62014-01-22 14:35:10 +00001165mp_obj_t mp_obj_str_builder_end(mp_obj_t o_in) {
Damien George5fa93b62014-01-22 14:35:10 +00001166 mp_obj_str_t *o = o_in;
1167 o->hash = qstr_compute_hash(o->data, o->len);
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001168 byte *p = (byte*)o->data;
1169 p[o->len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
Damien George5fa93b62014-01-22 14:35:10 +00001170 return o;
1171}
1172
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001173STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len) {
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001174 mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001175 o->base.type = type;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001176 o->len = len;
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001177 if (data) {
1178 o->hash = qstr_compute_hash(data, len);
1179 byte *p = m_new(byte, len + 1);
1180 o->data = p;
1181 memcpy(p, data, len * sizeof(byte));
1182 p[len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
1183 }
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001184 return o;
1185}
1186
Damien George5fa93b62014-01-22 14:35:10 +00001187mp_obj_t mp_obj_new_str(const byte* data, uint len, bool make_qstr_if_not_already) {
1188 qstr q = qstr_find_strn(data, len);
1189 if (q != MP_QSTR_NULL) {
1190 // qstr with this data already exists
1191 return MP_OBJ_NEW_QSTR(q);
1192 } else if (make_qstr_if_not_already) {
1193 // no existing qstr, make a new one
1194 return MP_OBJ_NEW_QSTR(qstr_from_strn((const char*)data, len));
1195 } else {
1196 // no existing qstr, don't make one
Damien George3e1a5c12014-03-29 13:43:38 +00001197 return str_new(&mp_type_str, data, len);
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +02001198 }
Damien George5fa93b62014-01-22 14:35:10 +00001199}
1200
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001201mp_obj_t mp_obj_new_bytes(const byte* data, uint len) {
Damien George3e1a5c12014-03-29 13:43:38 +00001202 return str_new(&mp_type_bytes, data, len);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001203}
1204
Damien George5fa93b62014-01-22 14:35:10 +00001205bool mp_obj_str_equal(mp_obj_t s1, mp_obj_t s2) {
1206 if (MP_OBJ_IS_QSTR(s1) && MP_OBJ_IS_QSTR(s2)) {
1207 return s1 == s2;
1208 } else {
1209 GET_STR_HASH(s1, h1);
1210 GET_STR_HASH(s2, h2);
1211 if (h1 != h2) {
1212 return false;
1213 }
1214 GET_STR_DATA_LEN(s1, d1, l1);
1215 GET_STR_DATA_LEN(s2, d2, l2);
1216 if (l1 != l2) {
1217 return false;
1218 }
Damien George1e708fe2014-01-23 18:27:51 +00001219 return memcmp(d1, d2, l1) == 0;
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +02001220 }
Damien George5fa93b62014-01-22 14:35:10 +00001221}
1222
Damien Georgeb829b5c2014-01-25 13:51:19 +00001223void bad_implicit_conversion(mp_obj_t self_in) __attribute__((noreturn));
1224void bad_implicit_conversion(mp_obj_t self_in) {
Damien Georgec5966122014-02-15 16:10:44 +00001225 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 +00001226}
1227
Damien George5fa93b62014-01-22 14:35:10 +00001228uint mp_obj_str_get_hash(mp_obj_t self_in) {
1229 if (MP_OBJ_IS_STR(self_in)) {
1230 GET_STR_HASH(self_in, h);
1231 return h;
1232 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001233 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001234 }
1235}
1236
1237uint mp_obj_str_get_len(mp_obj_t self_in) {
1238 if (MP_OBJ_IS_STR(self_in)) {
1239 GET_STR_LEN(self_in, l);
1240 return l;
1241 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001242 bad_implicit_conversion(self_in);
1243 }
1244}
1245
1246// use this if you will anyway convert the string to a qstr
1247// will be more efficient for the case where it's already a qstr
1248qstr mp_obj_str_get_qstr(mp_obj_t self_in) {
1249 if (MP_OBJ_IS_QSTR(self_in)) {
1250 return MP_OBJ_QSTR_VALUE(self_in);
Damien George3e1a5c12014-03-29 13:43:38 +00001251 } else if (MP_OBJ_IS_TYPE(self_in, &mp_type_str)) {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001252 mp_obj_str_t *self = self_in;
1253 return qstr_from_strn((char*)self->data, self->len);
1254 } else {
1255 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001256 }
1257}
1258
1259// only use this function if you need the str data to be zero terminated
1260// at the moment all strings are zero terminated to help with C ASCIIZ compatibility
1261const char *mp_obj_str_get_str(mp_obj_t self_in) {
1262 if (MP_OBJ_IS_STR(self_in)) {
1263 GET_STR_DATA_LEN(self_in, s, l);
1264 (void)l; // len unused
1265 return (const char*)s;
1266 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001267 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001268 }
1269}
1270
Damien George698ec212014-02-08 18:17:23 +00001271const char *mp_obj_str_get_data(mp_obj_t self_in, uint *len) {
Damien George5fa93b62014-01-22 14:35:10 +00001272 if (MP_OBJ_IS_STR(self_in)) {
1273 GET_STR_DATA_LEN(self_in, s, l);
1274 *len = l;
Damien George698ec212014-02-08 18:17:23 +00001275 return (const char*)s;
Damien George5fa93b62014-01-22 14:35:10 +00001276 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001277 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001278 }
Damiend99b0522013-12-21 18:17:45 +00001279}
xyb8cfc9f02014-01-05 18:47:51 +08001280
1281/******************************************************************************/
1282/* str iterator */
1283
1284typedef struct _mp_obj_str_it_t {
1285 mp_obj_base_t base;
Damien George5fa93b62014-01-22 14:35:10 +00001286 mp_obj_t str;
xyb8cfc9f02014-01-05 18:47:51 +08001287 machine_uint_t cur;
1288} mp_obj_str_it_t;
1289
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001290STATIC mp_obj_t str_it_iternext(mp_obj_t self_in) {
xyb8cfc9f02014-01-05 18:47:51 +08001291 mp_obj_str_it_t *self = self_in;
Damien George5fa93b62014-01-22 14:35:10 +00001292 GET_STR_DATA_LEN(self->str, str, len);
1293 if (self->cur < len) {
1294 mp_obj_t o_out = mp_obj_new_str(str + self->cur, 1, true);
xyb8cfc9f02014-01-05 18:47:51 +08001295 self->cur += 1;
1296 return o_out;
1297 } else {
Damien George66eaf842014-03-26 19:27:58 +00001298 return MP_OBJ_NULL;
xyb8cfc9f02014-01-05 18:47:51 +08001299 }
1300}
1301
Damien George3e1a5c12014-03-29 13:43:38 +00001302STATIC const mp_obj_type_t mp_type_str_it = {
Damien Georgec5966122014-02-15 16:10:44 +00001303 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001304 .name = MP_QSTR_iterator,
Paul Sokolovskyf7eaf602014-03-30 22:00:12 +03001305 .getiter = mp_identity,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001306 .iternext = str_it_iternext,
xyb8cfc9f02014-01-05 18:47:51 +08001307};
1308
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001309STATIC mp_obj_t bytes_it_iternext(mp_obj_t self_in) {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001310 mp_obj_str_it_t *self = self_in;
1311 GET_STR_DATA_LEN(self->str, str, len);
1312 if (self->cur < len) {
Damien George7c9c6672014-01-25 00:17:36 +00001313 mp_obj_t o_out = MP_OBJ_NEW_SMALL_INT((mp_small_int_t)str[self->cur]);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001314 self->cur += 1;
1315 return o_out;
1316 } else {
Damien George66eaf842014-03-26 19:27:58 +00001317 return MP_OBJ_NULL;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001318 }
1319}
1320
Damien George3e1a5c12014-03-29 13:43:38 +00001321STATIC const mp_obj_type_t mp_type_bytes_it = {
Damien Georgec5966122014-02-15 16:10:44 +00001322 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001323 .name = MP_QSTR_iterator,
Paul Sokolovskyf7eaf602014-03-30 22:00:12 +03001324 .getiter = mp_identity,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001325 .iternext = bytes_it_iternext,
1326};
1327
1328mp_obj_t mp_obj_new_str_iterator(mp_obj_t str) {
xyb8cfc9f02014-01-05 18:47:51 +08001329 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
Damien George3e1a5c12014-03-29 13:43:38 +00001330 o->base.type = &mp_type_str_it;
xyb8cfc9f02014-01-05 18:47:51 +08001331 o->str = str;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001332 o->cur = 0;
1333 return o;
1334}
1335
1336mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str) {
1337 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
Damien George3e1a5c12014-03-29 13:43:38 +00001338 o->base.type = &mp_type_bytes_it;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001339 o->str = str;
1340 o->cur = 0;
xyb8cfc9f02014-01-05 18:47:51 +08001341 return o;
1342}