blob: d0dee04e52ee3d468f471888667add3790cf66aa [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 }
Dave Hylands22fe4d72014-04-02 12:07:31 -0700791 }
792 // NOTE: no else here. We need the e, f, g etc formats for integer
793 // arguments (from above if) to take this if.
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700794#if MICROPY_ENABLE_FLOAT
Dave Hylands22fe4d72014-04-02 12:07:31 -0700795 if (arg_looks_numeric(arg)) {
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700796 if (!type) {
797
798 // Even though the docs say that an unspecified type is the same
799 // as 'g', there is one subtle difference, when the exponent
800 // is one less than the precision.
801 //
802 // '{:10.1}'.format(0.0) ==> '0e+00'
803 // '{:10.1g}'.format(0.0) ==> '0'
804 //
805 // TODO: Figure out how to deal with this.
806 //
807 // A proper solution would involve adding a special flag
808 // or something to format_float, and create a format_double
809 // to deal with doubles. In order to fix this when using
810 // sprintf, we'd need to use the e format and tweak the
811 // returned result to strip trailing zeros like the g format
812 // does.
813 //
814 // {:10.3} and {:10.2e} with 1.23e2 both produce 1.23e+02
815 // but with 1.e2 you get 1e+02 and 1.00e+02
816 //
817 // Stripping the trailing 0's (like g) does would make the
818 // e format give us the right format.
819 //
820 // CPython sources say:
821 // Omitted type specifier. Behaves in the same way as repr(x)
822 // and str(x) if no precision is given, else like 'g', but with
823 // at least one digit after the decimal point. */
824
825 type = 'g';
826 }
827 if (type == 'n') {
828 type = 'g';
829 }
830
831 flags |= PF_FLAG_PAD_NAN_INF; // '{:06e}'.format(float('-inf')) should give '-00inf'
832 switch (type) {
833 case 'e':
834 case 'E':
835 case 'f':
836 case 'F':
837 case 'g':
838 case 'G':
839 pfenv_print_float(&pfenv_vstr, mp_obj_get_float(arg), type, flags, fill, width, precision);
840 break;
841
842 case '%':
843 flags |= PF_FLAG_ADD_PERCENT;
844 pfenv_print_float(&pfenv_vstr, mp_obj_get_float(arg) * 100.0F, 'f', flags, fill, width, precision);
845 break;
846
847 default:
848 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError,
849 "Unknown format code '%c' for object of type 'float'",
850 type, mp_obj_get_type_str(arg)));
851 }
Dave Hylands22fe4d72014-04-02 12:07:31 -0700852 } else
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700853#endif
Damien George70f33cd2014-04-02 17:06:05 +0100854
Dave Hylands22fe4d72014-04-02 12:07:31 -0700855 {
Damien George70f33cd2014-04-02 17:06:05 +0100856 // arg doesn't look like a number
857
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700858 if (align == '=') {
859 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "'=' alignment not allowed in string format specifier"));
860 }
Damien George70f33cd2014-04-02 17:06:05 +0100861
Dave Hylandsbaf6f142014-03-30 21:06:50 -0700862 switch (type) {
863 case '\0':
864 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, arg, PRINT_STR);
865 break;
866
867 case 's':
868 {
869 uint len;
870 const char *s = mp_obj_str_get_data(arg, &len);
871 if (precision < 0) {
872 precision = len;
873 }
874 if (len > precision) {
875 len = precision;
876 }
877 pfenv_print_strn(&pfenv_vstr, s, len, flags, fill, width);
878 break;
879 }
880
881 default:
882 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_ValueError,
883 "Unknown format code '%c' for object of type 'str'",
884 type, mp_obj_get_type_str(arg)));
885 }
Damiend99b0522013-12-21 18:17:45 +0000886 }
887 }
888
Damien George5fa93b62014-01-22 14:35:10 +0000889 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
890 vstr_free(vstr);
891 return s;
Damiend99b0522013-12-21 18:17:45 +0000892}
893
Paul Sokolovsky4db727a2014-03-31 21:18:28 +0300894STATIC mp_obj_t str_modulo_format(mp_obj_t pattern, uint n_args, const mp_obj_t *args) {
895 assert(MP_OBJ_IS_STR(pattern));
896
897 GET_STR_DATA_LEN(pattern, str, len);
898 int arg_i = 0;
899 vstr_t *vstr = vstr_new();
900 for (const byte *top = str + len; str < top; str++) {
901 if (*str == '%') {
902 if (++str >= top) {
903 break;
904 }
905 if (*str == '%') {
906 vstr_add_char(vstr, '%');
907 } else {
908 if (arg_i >= n_args) {
909 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "not enough arguments for format string"));
910 }
911 switch (*str) {
912 case 's':
913 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR);
914 break;
915 case 'r':
916 mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_REPR);
917 break;
918 }
919 arg_i++;
920 }
921 } else {
922 vstr_add_char(vstr, *str);
923 }
924 }
925
926 if (arg_i != n_args) {
927 nlr_jump(mp_obj_new_exception_msg(&mp_type_TypeError, "not all arguments converted during string formatting"));
928 }
929
930 mp_obj_t s = mp_obj_new_str((byte*)vstr->buf, vstr->len, false);
931 vstr_free(vstr);
932 return s;
933}
934
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +0200935STATIC mp_obj_t str_replace(uint n_args, const mp_obj_t *args) {
xbe480c15a2014-01-30 22:17:30 -0800936 assert(MP_OBJ_IS_STR(args[0]));
937 assert(MP_OBJ_IS_STR(args[1]));
938 assert(MP_OBJ_IS_STR(args[2]));
939
Damien George94f68302014-01-31 23:45:12 +0000940 machine_int_t max_rep = 0;
xbe480c15a2014-01-30 22:17:30 -0800941 if (n_args == 4) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200942 assert(MP_OBJ_IS_SMALL_INT(args[3]));
943 max_rep = MP_OBJ_SMALL_INT_VALUE(args[3]);
944 if (max_rep == 0) {
945 return args[0];
946 } else if (max_rep < 0) {
947 max_rep = 0;
948 }
xbe480c15a2014-01-30 22:17:30 -0800949 }
Damien George94f68302014-01-31 23:45:12 +0000950
951 // if max_rep is still 0 by this point we will need to do all possible replacements
xbe480c15a2014-01-30 22:17:30 -0800952
953 GET_STR_DATA_LEN(args[0], str, str_len);
954 GET_STR_DATA_LEN(args[1], old, old_len);
955 GET_STR_DATA_LEN(args[2], new, new_len);
Damien George94f68302014-01-31 23:45:12 +0000956
957 // old won't exist in str if it's longer, so nothing to replace
xbe480c15a2014-01-30 22:17:30 -0800958 if (old_len > str_len) {
Paul Sokolovsky4e246082014-02-11 15:29:55 +0200959 return args[0];
xbe480c15a2014-01-30 22:17:30 -0800960 }
961
Damien George94f68302014-01-31 23:45:12 +0000962 // data for the replaced string
963 byte *data = NULL;
964 mp_obj_t replaced_str = MP_OBJ_NULL;
xbe480c15a2014-01-30 22:17:30 -0800965
Damien George94f68302014-01-31 23:45:12 +0000966 // do 2 passes over the string:
967 // first pass computes the required length of the replaced string
968 // second pass does the replacements
969 for (;;) {
970 machine_uint_t replaced_str_index = 0;
971 machine_uint_t num_replacements_done = 0;
972 const byte *old_occurrence;
973 const byte *offset_ptr = str;
974 machine_uint_t offset_num = 0;
xbe17a5a832014-03-23 23:31:58 -0700975 while ((old_occurrence = find_subbytes(offset_ptr, str_len - offset_num, old, old_len, 1)) != NULL) {
Damien George94f68302014-01-31 23:45:12 +0000976 // copy from just after end of last occurrence of to-be-replaced string to right before start of next occurrence
977 if (data != NULL) {
978 memcpy(data + replaced_str_index, offset_ptr, old_occurrence - offset_ptr);
979 }
980 replaced_str_index += old_occurrence - offset_ptr;
981 // copy the replacement string
982 if (data != NULL) {
983 memcpy(data + replaced_str_index, new, new_len);
984 }
985 replaced_str_index += new_len;
986 offset_ptr = old_occurrence + old_len;
987 offset_num = offset_ptr - str;
xbe480c15a2014-01-30 22:17:30 -0800988
Damien George94f68302014-01-31 23:45:12 +0000989 num_replacements_done++;
990 if (max_rep != 0 && num_replacements_done == max_rep){
991 break;
992 }
993 }
994
995 // copy from just after end of last occurrence of to-be-replaced string to end of old string
996 if (data != NULL) {
997 memcpy(data + replaced_str_index, offset_ptr, str_len - offset_num);
998 }
999 replaced_str_index += str_len - offset_num;
1000
1001 if (data == NULL) {
1002 // first pass
1003 if (num_replacements_done == 0) {
1004 // no substr found, return original string
1005 return args[0];
1006 } else {
1007 // substr found, allocate new string
1008 replaced_str = mp_obj_str_builder_start(mp_obj_get_type(args[0]), replaced_str_index, &data);
1009 }
1010 } else {
1011 // second pass, we are done
1012 break;
1013 }
xbe480c15a2014-01-30 22:17:30 -08001014 }
Damien George94f68302014-01-31 23:45:12 +00001015
xbe480c15a2014-01-30 22:17:30 -08001016 return mp_obj_str_builder_end(replaced_str);
1017}
1018
xbe9e1e8cd2014-03-12 22:57:16 -07001019STATIC mp_obj_t str_count(uint n_args, const mp_obj_t *args) {
1020 assert(2 <= n_args && n_args <= 4);
1021 assert(MP_OBJ_IS_STR(args[0]));
1022 assert(MP_OBJ_IS_STR(args[1]));
1023
1024 GET_STR_DATA_LEN(args[0], haystack, haystack_len);
1025 GET_STR_DATA_LEN(args[1], needle, needle_len);
1026
Damien George536dde22014-03-13 22:07:55 +00001027 machine_uint_t start = 0;
1028 machine_uint_t end = haystack_len;
xbe9e1e8cd2014-03-12 22:57:16 -07001029 if (n_args >= 3 && args[2] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +00001030 start = mp_get_index(&mp_type_str, haystack_len, args[2], true);
xbe9e1e8cd2014-03-12 22:57:16 -07001031 }
1032 if (n_args >= 4 && args[3] != mp_const_none) {
Damien George3e1a5c12014-03-29 13:43:38 +00001033 end = mp_get_index(&mp_type_str, haystack_len, args[3], true);
xbe9e1e8cd2014-03-12 22:57:16 -07001034 }
1035
Damien George536dde22014-03-13 22:07:55 +00001036 // if needle_len is zero then we count each gap between characters as an occurrence
1037 if (needle_len == 0) {
1038 return MP_OBJ_NEW_SMALL_INT(end - start + 1);
xbe9e1e8cd2014-03-12 22:57:16 -07001039 }
1040
Damien George536dde22014-03-13 22:07:55 +00001041 // count the occurrences
1042 machine_int_t num_occurrences = 0;
xbec5d70ba2014-03-13 00:29:15 -07001043 for (machine_uint_t haystack_index = start; haystack_index + needle_len <= end; haystack_index++) {
1044 if (memcmp(&haystack[haystack_index], needle, needle_len) == 0) {
1045 num_occurrences++;
1046 haystack_index += needle_len - 1;
1047 }
xbe9e1e8cd2014-03-12 22:57:16 -07001048 }
1049
1050 return MP_OBJ_NEW_SMALL_INT(num_occurrences);
1051}
1052
Damien Georgeb035db32014-03-21 20:39:40 +00001053STATIC mp_obj_t str_partitioner(mp_obj_t self_in, mp_obj_t arg, machine_int_t direction) {
xbe613a8e32014-03-18 00:06:29 -07001054 assert(MP_OBJ_IS_STR(self_in));
1055 if (!MP_OBJ_IS_STR(arg)) {
1056 nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError,
1057 "Can't convert '%s' object to str implicitly", mp_obj_get_type_str(arg)));
1058 }
Damien Georgeb035db32014-03-21 20:39:40 +00001059
xbe613a8e32014-03-18 00:06:29 -07001060 GET_STR_DATA_LEN(self_in, str, str_len);
1061 GET_STR_DATA_LEN(arg, sep, sep_len);
1062
1063 if (sep_len == 0) {
1064 nlr_jump(mp_obj_new_exception_msg(&mp_type_ValueError, "empty separator"));
1065 }
Damien Georgeb035db32014-03-21 20:39:40 +00001066
1067 mp_obj_t result[] = {MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_), MP_OBJ_NEW_QSTR(MP_QSTR_)};
1068
1069 if (direction > 0) {
1070 result[0] = self_in;
xbe0a6894c2014-03-21 01:12:26 -07001071 } else {
Damien Georgeb035db32014-03-21 20:39:40 +00001072 result[2] = self_in;
xbe0a6894c2014-03-21 01:12:26 -07001073 }
xbe613a8e32014-03-18 00:06:29 -07001074
xbe17a5a832014-03-23 23:31:58 -07001075 const byte *position_ptr = find_subbytes(str, str_len, sep, sep_len, direction);
1076 if (position_ptr != NULL) {
1077 machine_uint_t position = position_ptr - str;
1078 result[0] = mp_obj_new_str(str, position, false);
1079 result[1] = arg;
1080 result[2] = mp_obj_new_str(str + position + sep_len, str_len - position - sep_len, false);
xbe613a8e32014-03-18 00:06:29 -07001081 }
Damien Georgeb035db32014-03-21 20:39:40 +00001082
xbe0a6894c2014-03-21 01:12:26 -07001083 return mp_obj_new_tuple(3, result);
xbe613a8e32014-03-18 00:06:29 -07001084}
1085
Damien Georgeb035db32014-03-21 20:39:40 +00001086STATIC mp_obj_t str_partition(mp_obj_t self_in, mp_obj_t arg) {
1087 return str_partitioner(self_in, arg, 1);
xbe0a6894c2014-03-21 01:12:26 -07001088}
xbe4504ea82014-03-19 00:46:14 -07001089
Damien Georgeb035db32014-03-21 20:39:40 +00001090STATIC mp_obj_t str_rpartition(mp_obj_t self_in, mp_obj_t arg) {
1091 return str_partitioner(self_in, arg, -1);
xbe4504ea82014-03-19 00:46:14 -07001092}
1093
Damien George2da98302014-03-09 19:58:18 +00001094STATIC machine_int_t str_get_buffer(mp_obj_t self_in, buffer_info_t *bufinfo, int flags) {
1095 if (flags == BUFFER_READ) {
1096 GET_STR_DATA_LEN(self_in, str_data, str_len);
1097 bufinfo->buf = (void*)str_data;
1098 bufinfo->len = str_len;
1099 return 0;
1100 } else {
1101 // can't write to a string
1102 bufinfo->buf = NULL;
1103 bufinfo->len = 0;
1104 return 1;
1105 }
1106}
1107
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001108STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
xbe17a5a832014-03-23 23:31:58 -07001109STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_rfind_obj, 2, 4, str_rfind);
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001110STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
1111STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_split_obj, 1, 3, str_split);
1112STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_startswith_obj, str_startswith);
1113STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
1114STATIC MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);
1115STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_replace_obj, 3, 4, str_replace);
xbe9e1e8cd2014-03-12 22:57:16 -07001116STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_count_obj, 2, 4, str_count);
xbe613a8e32014-03-18 00:06:29 -07001117STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_partition_obj, str_partition);
xbe4504ea82014-03-19 00:46:14 -07001118STATIC MP_DEFINE_CONST_FUN_OBJ_2(str_rpartition_obj, str_rpartition);
Damiend99b0522013-12-21 18:17:45 +00001119
Damien George9b196cd2014-03-26 21:47:19 +00001120STATIC const mp_map_elem_t str_locals_dict_table[] = {
1121 { MP_OBJ_NEW_QSTR(MP_QSTR_find), (mp_obj_t)&str_find_obj },
1122 { MP_OBJ_NEW_QSTR(MP_QSTR_rfind), (mp_obj_t)&str_rfind_obj },
1123 { MP_OBJ_NEW_QSTR(MP_QSTR_join), (mp_obj_t)&str_join_obj },
1124 { MP_OBJ_NEW_QSTR(MP_QSTR_split), (mp_obj_t)&str_split_obj },
1125 { MP_OBJ_NEW_QSTR(MP_QSTR_startswith), (mp_obj_t)&str_startswith_obj },
1126 { MP_OBJ_NEW_QSTR(MP_QSTR_strip), (mp_obj_t)&str_strip_obj },
1127 { MP_OBJ_NEW_QSTR(MP_QSTR_format), (mp_obj_t)&str_format_obj },
1128 { MP_OBJ_NEW_QSTR(MP_QSTR_replace), (mp_obj_t)&str_replace_obj },
1129 { MP_OBJ_NEW_QSTR(MP_QSTR_count), (mp_obj_t)&str_count_obj },
1130 { MP_OBJ_NEW_QSTR(MP_QSTR_partition), (mp_obj_t)&str_partition_obj },
1131 { MP_OBJ_NEW_QSTR(MP_QSTR_rpartition), (mp_obj_t)&str_rpartition_obj },
ian-v7a16fad2014-01-06 09:52:29 -08001132};
Damien George97209d32014-01-07 15:58:30 +00001133
Damien George9b196cd2014-03-26 21:47:19 +00001134STATIC MP_DEFINE_CONST_DICT(str_locals_dict, str_locals_dict_table);
1135
Damien George3e1a5c12014-03-29 13:43:38 +00001136const mp_obj_type_t mp_type_str = {
Damien Georgec5966122014-02-15 16:10:44 +00001137 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001138 .name = MP_QSTR_str,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001139 .print = str_print,
Paul Sokolovskybe020c22014-03-21 11:39:01 +02001140 .make_new = str_make_new,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001141 .binary_op = str_binary_op,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001142 .getiter = mp_obj_new_str_iterator,
Damien George2da98302014-03-09 19:58:18 +00001143 .buffer_p = { .get_buffer = str_get_buffer },
Damien George9b196cd2014-03-26 21:47:19 +00001144 .locals_dict = (mp_obj_t)&str_locals_dict,
Damiend99b0522013-12-21 18:17:45 +00001145};
1146
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001147// Reuses most of methods from str
Damien George3e1a5c12014-03-29 13:43:38 +00001148const mp_obj_type_t mp_type_bytes = {
Damien Georgec5966122014-02-15 16:10:44 +00001149 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001150 .name = MP_QSTR_bytes,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001151 .print = str_print,
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001152 .make_new = bytes_make_new,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001153 .binary_op = str_binary_op,
1154 .getiter = mp_obj_new_bytes_iterator,
Damien George9b196cd2014-03-26 21:47:19 +00001155 .locals_dict = (mp_obj_t)&str_locals_dict,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001156};
1157
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001158// the zero-length bytes
Damien George3e1a5c12014-03-29 13:43:38 +00001159STATIC const mp_obj_str_t empty_bytes_obj = {{&mp_type_bytes}, 0, 0, NULL};
Paul Sokolovsky1ecea7c2014-03-21 23:46:59 +02001160const mp_obj_t mp_const_empty_bytes = (mp_obj_t)&empty_bytes_obj;
1161
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001162mp_obj_t mp_obj_str_builder_start(const mp_obj_type_t *type, uint len, byte **data) {
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001163 mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001164 o->base.type = type;
Damien George5fa93b62014-01-22 14:35:10 +00001165 o->len = len;
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001166 byte *p = m_new(byte, len + 1);
1167 o->data = p;
1168 *data = p;
Damiend99b0522013-12-21 18:17:45 +00001169 return o;
1170}
1171
Damien George5fa93b62014-01-22 14:35:10 +00001172mp_obj_t mp_obj_str_builder_end(mp_obj_t o_in) {
Damien George5fa93b62014-01-22 14:35:10 +00001173 mp_obj_str_t *o = o_in;
1174 o->hash = qstr_compute_hash(o->data, o->len);
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001175 byte *p = (byte*)o->data;
1176 p[o->len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
Damien George5fa93b62014-01-22 14:35:10 +00001177 return o;
1178}
1179
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001180STATIC mp_obj_t str_new(const mp_obj_type_t *type, const byte* data, uint len) {
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001181 mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001182 o->base.type = type;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001183 o->len = len;
Paul Sokolovsky5972b4c2014-03-20 16:47:44 +02001184 if (data) {
1185 o->hash = qstr_compute_hash(data, len);
1186 byte *p = m_new(byte, len + 1);
1187 o->data = p;
1188 memcpy(p, data, len * sizeof(byte));
1189 p[len] = '\0'; // for now we add null for compatibility with C ASCIIZ strings
1190 }
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001191 return o;
1192}
1193
Damien George5fa93b62014-01-22 14:35:10 +00001194mp_obj_t mp_obj_new_str(const byte* data, uint len, bool make_qstr_if_not_already) {
1195 qstr q = qstr_find_strn(data, len);
1196 if (q != MP_QSTR_NULL) {
1197 // qstr with this data already exists
1198 return MP_OBJ_NEW_QSTR(q);
1199 } else if (make_qstr_if_not_already) {
1200 // no existing qstr, make a new one
1201 return MP_OBJ_NEW_QSTR(qstr_from_strn((const char*)data, len));
1202 } else {
1203 // no existing qstr, don't make one
Damien George3e1a5c12014-03-29 13:43:38 +00001204 return str_new(&mp_type_str, data, len);
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +02001205 }
Damien George5fa93b62014-01-22 14:35:10 +00001206}
1207
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001208mp_obj_t mp_obj_new_bytes(const byte* data, uint len) {
Damien George3e1a5c12014-03-29 13:43:38 +00001209 return str_new(&mp_type_bytes, data, len);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001210}
1211
Damien George5fa93b62014-01-22 14:35:10 +00001212bool mp_obj_str_equal(mp_obj_t s1, mp_obj_t s2) {
1213 if (MP_OBJ_IS_QSTR(s1) && MP_OBJ_IS_QSTR(s2)) {
1214 return s1 == s2;
1215 } else {
1216 GET_STR_HASH(s1, h1);
1217 GET_STR_HASH(s2, h2);
1218 if (h1 != h2) {
1219 return false;
1220 }
1221 GET_STR_DATA_LEN(s1, d1, l1);
1222 GET_STR_DATA_LEN(s2, d2, l2);
1223 if (l1 != l2) {
1224 return false;
1225 }
Damien George1e708fe2014-01-23 18:27:51 +00001226 return memcmp(d1, d2, l1) == 0;
Paul Sokolovsky8965a5e2014-01-20 23:33:19 +02001227 }
Damien George5fa93b62014-01-22 14:35:10 +00001228}
1229
Damien Georgeb829b5c2014-01-25 13:51:19 +00001230void bad_implicit_conversion(mp_obj_t self_in) __attribute__((noreturn));
1231void bad_implicit_conversion(mp_obj_t self_in) {
Damien Georgec5966122014-02-15 16:10:44 +00001232 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 +00001233}
1234
Damien George5fa93b62014-01-22 14:35:10 +00001235uint mp_obj_str_get_hash(mp_obj_t self_in) {
1236 if (MP_OBJ_IS_STR(self_in)) {
1237 GET_STR_HASH(self_in, h);
1238 return h;
1239 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001240 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001241 }
1242}
1243
1244uint mp_obj_str_get_len(mp_obj_t self_in) {
1245 if (MP_OBJ_IS_STR(self_in)) {
1246 GET_STR_LEN(self_in, l);
1247 return l;
1248 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001249 bad_implicit_conversion(self_in);
1250 }
1251}
1252
1253// use this if you will anyway convert the string to a qstr
1254// will be more efficient for the case where it's already a qstr
1255qstr mp_obj_str_get_qstr(mp_obj_t self_in) {
1256 if (MP_OBJ_IS_QSTR(self_in)) {
1257 return MP_OBJ_QSTR_VALUE(self_in);
Damien George3e1a5c12014-03-29 13:43:38 +00001258 } else if (MP_OBJ_IS_TYPE(self_in, &mp_type_str)) {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001259 mp_obj_str_t *self = self_in;
1260 return qstr_from_strn((char*)self->data, self->len);
1261 } else {
1262 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001263 }
1264}
1265
1266// only use this function if you need the str data to be zero terminated
1267// at the moment all strings are zero terminated to help with C ASCIIZ compatibility
1268const char *mp_obj_str_get_str(mp_obj_t self_in) {
1269 if (MP_OBJ_IS_STR(self_in)) {
1270 GET_STR_DATA_LEN(self_in, s, l);
1271 (void)l; // len unused
1272 return (const char*)s;
1273 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001274 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001275 }
1276}
1277
Damien George698ec212014-02-08 18:17:23 +00001278const char *mp_obj_str_get_data(mp_obj_t self_in, uint *len) {
Damien George5fa93b62014-01-22 14:35:10 +00001279 if (MP_OBJ_IS_STR(self_in)) {
1280 GET_STR_DATA_LEN(self_in, s, l);
1281 *len = l;
Damien George698ec212014-02-08 18:17:23 +00001282 return (const char*)s;
Damien George5fa93b62014-01-22 14:35:10 +00001283 } else {
Damien Georgeb829b5c2014-01-25 13:51:19 +00001284 bad_implicit_conversion(self_in);
Damien George5fa93b62014-01-22 14:35:10 +00001285 }
Damiend99b0522013-12-21 18:17:45 +00001286}
xyb8cfc9f02014-01-05 18:47:51 +08001287
1288/******************************************************************************/
1289/* str iterator */
1290
1291typedef struct _mp_obj_str_it_t {
1292 mp_obj_base_t base;
Damien George5fa93b62014-01-22 14:35:10 +00001293 mp_obj_t str;
xyb8cfc9f02014-01-05 18:47:51 +08001294 machine_uint_t cur;
1295} mp_obj_str_it_t;
1296
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001297STATIC mp_obj_t str_it_iternext(mp_obj_t self_in) {
xyb8cfc9f02014-01-05 18:47:51 +08001298 mp_obj_str_it_t *self = self_in;
Damien George5fa93b62014-01-22 14:35:10 +00001299 GET_STR_DATA_LEN(self->str, str, len);
1300 if (self->cur < len) {
1301 mp_obj_t o_out = mp_obj_new_str(str + self->cur, 1, true);
xyb8cfc9f02014-01-05 18:47:51 +08001302 self->cur += 1;
1303 return o_out;
1304 } else {
Damien George66eaf842014-03-26 19:27:58 +00001305 return MP_OBJ_NULL;
xyb8cfc9f02014-01-05 18:47:51 +08001306 }
1307}
1308
Damien George3e1a5c12014-03-29 13:43:38 +00001309STATIC const mp_obj_type_t mp_type_str_it = {
Damien Georgec5966122014-02-15 16:10:44 +00001310 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001311 .name = MP_QSTR_iterator,
Paul Sokolovskyf7eaf602014-03-30 22:00:12 +03001312 .getiter = mp_identity,
Paul Sokolovsky860ffb02014-01-05 22:34:09 +02001313 .iternext = str_it_iternext,
xyb8cfc9f02014-01-05 18:47:51 +08001314};
1315
Paul Sokolovskyd5df6cd2014-02-12 18:15:40 +02001316STATIC mp_obj_t bytes_it_iternext(mp_obj_t self_in) {
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001317 mp_obj_str_it_t *self = self_in;
1318 GET_STR_DATA_LEN(self->str, str, len);
1319 if (self->cur < len) {
Damien George7c9c6672014-01-25 00:17:36 +00001320 mp_obj_t o_out = MP_OBJ_NEW_SMALL_INT((mp_small_int_t)str[self->cur]);
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001321 self->cur += 1;
1322 return o_out;
1323 } else {
Damien George66eaf842014-03-26 19:27:58 +00001324 return MP_OBJ_NULL;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001325 }
1326}
1327
Damien George3e1a5c12014-03-29 13:43:38 +00001328STATIC const mp_obj_type_t mp_type_bytes_it = {
Damien Georgec5966122014-02-15 16:10:44 +00001329 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +00001330 .name = MP_QSTR_iterator,
Paul Sokolovskyf7eaf602014-03-30 22:00:12 +03001331 .getiter = mp_identity,
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001332 .iternext = bytes_it_iternext,
1333};
1334
1335mp_obj_t mp_obj_new_str_iterator(mp_obj_t str) {
xyb8cfc9f02014-01-05 18:47:51 +08001336 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
Damien George3e1a5c12014-03-29 13:43:38 +00001337 o->base.type = &mp_type_str_it;
xyb8cfc9f02014-01-05 18:47:51 +08001338 o->str = str;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001339 o->cur = 0;
1340 return o;
1341}
1342
1343mp_obj_t mp_obj_new_bytes_iterator(mp_obj_t str) {
1344 mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
Damien George3e1a5c12014-03-29 13:43:38 +00001345 o->base.type = &mp_type_bytes_it;
Paul Sokolovsky91fb1c92014-01-24 22:50:40 +02001346 o->str = str;
1347 o->cur = 0;
xyb8cfc9f02014-01-05 18:47:51 +08001348 return o;
1349}