objstr.c 12 KB
Newer Older
1
2
3
4
5
6
7
8
9
#include <stdlib.h>
#include <stdint.h>
#include <stdarg.h>
#include <string.h>
#include <assert.h>

#include "nlr.h"
#include "misc.h"
#include "mpconfig.h"
10
#include "qstr.h"
11
12
13
14
15
16
17
18
19
#include "obj.h"
#include "runtime0.h"
#include "runtime.h"

typedef struct _mp_obj_str_t {
    mp_obj_base_t base;
    qstr qstr;
} mp_obj_str_t;

xyb's avatar
xyb committed
20
21
22
23
24
static mp_obj_t mp_obj_new_str_iterator(mp_obj_str_t *str, int cur);

/******************************************************************************/
/* str                                                                        */

25
void mp_obj_str_print_qstr(void (*print)(void *env, const char *fmt, ...), void *env, qstr q, mp_print_kind_t kind) {
26
    if (kind == PRINT_STR) {
27
        print(env, "%s", qstr_str(q));
28
29
    } else {
        // TODO need to escape chars etc
30
        print(env, "'%s'", qstr_str(q));
31
    }
32
33
}

34
35
36
37
38
void str_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in, mp_print_kind_t kind) {
    mp_obj_str_t *self = self_in;
    mp_obj_str_print_qstr(print, env, self->qstr, kind);
}

39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
// like strstr but with specified length and allows \0 bytes
// TODO replace with something more efficient/standard
static const byte *find_subbytes(const byte *haystack, uint hlen, const byte *needle, uint nlen) {
    if (hlen >= nlen) {
        for (uint i = 0; i <= hlen - nlen; i++) {
            bool found = true;
            for (uint j = 0; j < nlen; j++) {
                if (haystack[i + j] != needle[j]) {
                    found = false;
                    break;
                }
            }
            if (found) {
                return haystack + i;
            }
        }
    }
    return NULL;
}

59
60
mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
    mp_obj_str_t *lhs = lhs_in;
61
62
    uint lhs_len;
    const byte *lhs_data = qstr_data(lhs->qstr, &lhs_len);
63
64
    switch (op) {
        case RT_BINARY_OP_SUBSCR:
65
66
67
68
            // TODO: need predicate to check for int-like type (bools are such for example)
            // ["no", "yes"][1 == 2] is common idiom
            if (MP_OBJ_IS_SMALL_INT(rhs_in)) {
                // TODO: This implements byte string access for single index so far
69
                // TODO: Handle negative indexes.
70
                return mp_obj_new_int(lhs_data[mp_obj_get_int(rhs_in)]);
71
#if MICROPY_ENABLE_SLICE
72
            } else if (MP_OBJ_IS_TYPE(rhs_in, &slice_type)) {
73
                machine_int_t start, stop, step;
74
75
                mp_obj_slice_get(rhs_in, &start, &stop, &step);
                assert(step == 1);
76
                if (start < 0) {
77
                    start = lhs_len + start;
78
79
80
                    if (start < 0) {
                        start = 0;
                    }
81
82
                } else if (start > lhs_len) {
                    start = lhs_len;
83
84
                }
                if (stop <= 0) {
85
                    stop = lhs_len + stop;
86
87
88
89
                    // CPython returns empty string in such case
                    if (stop < 0) {
                        stop = start;
                    }
90
91
                } else if (stop > lhs_len) {
                    stop = lhs_len;
92
                }
93
                return mp_obj_new_str(qstr_from_strn((const char*)lhs_data + start, stop - start));
94
#endif
95
            } else {
96
97
                // Message doesn't match CPython, but we don't have so much bytes as they
                // to spend them on verbose wording
98
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "index must be int"));
99
            }
100
101
102
103
104

        case RT_BINARY_OP_ADD:
        case RT_BINARY_OP_INPLACE_ADD:
            if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
                // add 2 strings
105
106
107
108
109
110
111
112
                uint rhs_len;
                const byte *rhs_data = qstr_data(((mp_obj_str_t*)rhs_in)->qstr, &rhs_len);
                int alloc_len = lhs_len + rhs_len;
                byte *q_ptr;
                byte *val = qstr_build_start(alloc_len, &q_ptr);
                memcpy(val, lhs_data, lhs_len);
                memcpy(val + lhs_len, rhs_data, rhs_len);
                return mp_obj_new_str(qstr_build_end(q_ptr));
113
114
            }
            break;
115
116
117
118
        case RT_COMPARE_OP_IN:
        case RT_COMPARE_OP_NOT_IN:
            /* NOTE `a in b` is `b.__contains__(a)` */
            if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
119
120
121
122
                uint rhs_len;
                const byte *rhs_data = qstr_data(((mp_obj_str_t*)rhs_in)->qstr, &rhs_len);
                return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (find_subbytes(lhs_data, lhs_len, rhs_data, rhs_len) == NULL));
                return mp_const_false;
123
124
            }
            break;
125
126
127
128
129
    }

    return MP_OBJ_NULL; // op not supported
}

xyb's avatar
xyb committed
130
131
132
133
static mp_obj_t str_getiter(mp_obj_t o_in) {
    return mp_obj_new_str_iterator(o_in, 0);
}

134
135
136
137
mp_obj_t str_join(mp_obj_t self_in, mp_obj_t arg) {
    assert(MP_OBJ_IS_TYPE(self_in, &str_type));
    mp_obj_str_t *self = self_in;

138
139
140
141
142
    // get separation string
    const char *sep_str = qstr_str(self->qstr);
    size_t sep_len = strlen(sep_str);

    // process args
143
144
145
146
147
148
149
150
151
    uint seq_len;
    mp_obj_t *seq_items;
    if (MP_OBJ_IS_TYPE(arg, &tuple_type)) {
        mp_obj_tuple_get(arg, &seq_len, &seq_items);
    } else if (MP_OBJ_IS_TYPE(arg, &list_type)) {
        mp_obj_list_get(arg, &seq_len, &seq_items);
    } else {
        goto bad_arg;
    }
152
153
154

    // count required length
    int required_len = 0;
155
156
157
158
    for (int i = 0; i < seq_len; i++) {
        if (!MP_OBJ_IS_TYPE(seq_items[i], &str_type)) {
            goto bad_arg;
        }
159
160
161
        if (i > 0) {
            required_len += sep_len;
        }
162
163
164
165
        required_len += strlen(qstr_str(mp_obj_str_get(seq_items[i])));
    }

    // make joined string
166
167
    byte *q_ptr;
    byte *s_dest = qstr_build_start(required_len, &q_ptr);
168
169
    for (int i = 0; i < seq_len; i++) {
        if (i > 0) {
170
171
            memcpy(s_dest, sep_str, sep_len);
            s_dest += sep_len;
172
        }
173
174
        uint s2_len;
        const byte *s2 = qstr_data(mp_obj_str_get(seq_items[i]), &s2_len);
175
176
        memcpy(s_dest, s2, s2_len);
        s_dest += s2_len;
177
    }
178
    qstr q = qstr_build_end(q_ptr);
179
180

    // return joined string
181
    return mp_obj_new_str(q);
182
183

bad_arg:
184
    nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "?str.join expecting a list of str's"));
185
186
}

xbe's avatar
xbe committed
187
188
189
190
191
192
193
194
195
static bool chr_in_str(const char* const str, const size_t str_len, const char c) {
    for (size_t i = 0; i < str_len; i++) {
        if (str[i] == c) {
            return true;
        }
    }
    return false;
}

196
static mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
197
198
199
200
201
202
203
204
205
206
207
208
    assert(2 <= n_args && n_args <= 4);
    assert(MP_OBJ_IS_TYPE(args[0], &str_type));
    if (!MP_OBJ_IS_TYPE(args[1], &str_type)) {
        nlr_jump(mp_obj_new_exception_msg_1_arg(
                     MP_QSTR_TypeError,
                     "Can't convert '%s' object to str implicitly",
                     mp_obj_get_type_str(args[1])));
    }

    const char* haystack = qstr_str(((mp_obj_str_t*)args[0])->qstr);
    const char* needle = qstr_str(((mp_obj_str_t*)args[1])->qstr);

209
210
    size_t haystack_len = strlen(haystack);
    size_t needle_len = strlen(needle);
211
212
213
214
215
216
217
218
219
220
221
222

    size_t start = 0;
    size_t end = haystack_len;
    /* TODO use a non-exception-throwing mp_get_index */
    if (n_args >= 3 && args[2] != mp_const_none) {
        start = mp_get_index(&str_type, haystack_len, args[2]);
    }
    if (n_args >= 4 && args[3] != mp_const_none) {
        end = mp_get_index(&str_type, haystack_len, args[3]);
    }

    char *p = strstr(haystack + start, needle);
223
224
225
226
227
228
    if (p == NULL) {
        // not found
        return MP_OBJ_NEW_SMALL_INT(-1);
    } else {
        // found
        machine_int_t pos = p - haystack;
229
230
231
        if (pos + needle_len > end) {
            pos = -1;
        }
232
        return MP_OBJ_NEW_SMALL_INT(pos);
233
234
235
    }
}

236
mp_obj_t str_strip(uint n_args, const mp_obj_t *args) {
xbe's avatar
xbe committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
    assert(1 <= n_args && n_args <= 2);
    assert(MP_OBJ_IS_TYPE(args[0], &str_type));
    const char *chars_to_del;
    static const char whitespace[] = " \t\n\r\v\f";

    if (n_args == 1) {
        chars_to_del = whitespace;
    } else {
        assert(MP_OBJ_IS_TYPE(args[1], &str_type));
        mp_obj_str_t *chars_to_del_obj = args[1];
        chars_to_del = qstr_str(chars_to_del_obj->qstr);
    }

    const size_t chars_to_del_len = strlen(chars_to_del);
    mp_obj_str_t *self = args[0];
    const char *orig_str = qstr_str(self->qstr);
    const size_t orig_str_len = strlen(orig_str);

    size_t first_good_char_pos = 0;
    bool first_good_char_pos_set = false;
    size_t last_good_char_pos = 0;
    for (size_t i = 0; i < orig_str_len; i++) {
        if (!chr_in_str(chars_to_del, chars_to_del_len, orig_str[i])) {
            last_good_char_pos = i;
            if (!first_good_char_pos_set) {
                first_good_char_pos = i;
                first_good_char_pos_set = true;
            }
        }
    }

    if (first_good_char_pos == 0 && last_good_char_pos == 0) {
269
270
        //string is all whitespace, return ''
        return mp_obj_new_str(MP_QSTR_);
xbe's avatar
xbe committed
271
272
273
274
275
    }

    assert(last_good_char_pos >= first_good_char_pos);
    //+1 to accomodate the last character
    size_t stripped_len = last_good_char_pos - first_good_char_pos + 1;
276
    return mp_obj_new_str(qstr_from_strn(orig_str + first_good_char_pos, stripped_len));
xbe's avatar
xbe committed
277
278
}

279
mp_obj_t str_format(uint n_args, const mp_obj_t *args) {
280
281
282
283
284
285
286
287
288
289
290
    assert(MP_OBJ_IS_TYPE(args[0], &str_type));
    mp_obj_str_t *self = args[0];

    const char *str = qstr_str(self->qstr);
    int arg_i = 1;
    vstr_t *vstr = vstr_new();
    for (; *str; str++) {
        if (*str == '{') {
            str++;
            if (*str == '{') {
                vstr_add_char(vstr, '{');
291
292
            } else {
                while (*str != '}') str++;
293
                if (arg_i >= n_args) {
294
                    nlr_jump(mp_obj_new_exception_msg(MP_QSTR_IndexError, "tuple index out of range"));
295
                }
296
                // TODO: may be PRINT_REPR depending on formatting code
297
                mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR);
298
299
300
301
302
303
304
                arg_i++;
            }
        } else {
            vstr_add_char(vstr, *str);
        }
    }

305
    return mp_obj_new_str(qstr_from_strn_take(vstr->buf, vstr->alloc, vstr->len));
306
307
}

308
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
309
static MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
xbe's avatar
xbe committed
310
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
311
312
static MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);

ian-v's avatar
ian-v committed
313
static const mp_method_t str_type_methods[] = {
314
    { "find", &str_find_obj },
ian-v's avatar
ian-v committed
315
    { "join", &str_join_obj },
xbe's avatar
xbe committed
316
    { "strip", &str_strip_obj },
ian-v's avatar
ian-v committed
317
318
319
    { "format", &str_format_obj },
    { NULL, NULL }, // end-of-list sentinel
};
320

321
322
323
const mp_obj_type_t str_type = {
    { &mp_const_type },
    "str",
324
325
326
    .print = str_print,
    .binary_op = str_binary_op,
    .getiter = str_getiter,
ian-v's avatar
ian-v committed
327
    .methods = str_type_methods,
328
329
330
331
332
333
334
335
336
337
338
339
340
341
};

mp_obj_t mp_obj_new_str(qstr qstr) {
    mp_obj_str_t *o = m_new_obj(mp_obj_str_t);
    o->base.type = &str_type;
    o->qstr = qstr;
    return o;
}

qstr mp_obj_str_get(mp_obj_t self_in) {
    assert(MP_OBJ_IS_TYPE(self_in, &str_type));
    mp_obj_str_t *self = self_in;
    return self->qstr;
}
xyb's avatar
xyb committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355

/******************************************************************************/
/* str iterator                                                               */

typedef struct _mp_obj_str_it_t {
    mp_obj_base_t base;
    mp_obj_str_t *str;
    machine_uint_t cur;
} mp_obj_str_it_t;

mp_obj_t str_it_iternext(mp_obj_t self_in) {
    mp_obj_str_it_t *self = self_in;
    const char *str = qstr_str(self->str->qstr);
    if (self->cur < strlen(str)) {
356
        mp_obj_t o_out = mp_obj_new_str(qstr_from_strn(str + self->cur, 1));
xyb's avatar
xyb committed
357
358
359
360
361
362
363
364
365
366
        self->cur += 1;
        return o_out;
    } else {
        return mp_const_stop_iteration;
    }
}

static const mp_obj_type_t str_it_type = {
    { &mp_const_type },
    "str_iterator",
367
    .iternext = str_it_iternext,
xyb's avatar
xyb committed
368
369
370
371
372
373
374
375
376
};

mp_obj_t mp_obj_new_str_iterator(mp_obj_str_t *str, int cur) {
    mp_obj_str_it_t *o = m_new_obj(mp_obj_str_it_t);
    o->base.type = &str_it_type;
    o->str = str;
    o->cur = cur;
    return o;
}