objstr.c 11.9 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 "mpqstr.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
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;
    const char *lhs_str = qstr_str(lhs->qstr);
    switch (op) {
        case RT_BINARY_OP_SUBSCR:
44
45
46
            // 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)) {
47
48
                uint index = mp_get_index(lhs->base.type, strlen(lhs_str), rhs_in);
                return mp_obj_new_str(qstr_from_strn_copy(lhs_str + index, 1));
49
#if MICROPY_ENABLE_SLICE
50
            } else if (MP_OBJ_IS_TYPE(rhs_in, &slice_type)) {
51
                machine_int_t start, stop, step;
52
53
                mp_obj_slice_get(rhs_in, &start, &stop, &step);
                assert(step == 1);
54
55
56
                int len = strlen(lhs_str);
                if (start < 0) {
                    start = len + start;
57
58
59
60
61
                    if (start < 0) {
                        start = 0;
                    }
                } else if (start > len) {
                    start = len;
62
63
64
                }
                if (stop <= 0) {
                    stop = len + stop;
65
66
67
68
69
70
                    // CPython returns empty string in such case
                    if (stop < 0) {
                        stop = start;
                    }
                } else if (stop > len) {
                    stop = len;
71
                }
72
                return mp_obj_new_str(qstr_from_strn_copy(lhs_str + start, stop - start));
73
#endif
74
            } else {
75
76
                // Message doesn't match CPython, but we don't have so much bytes as they
                // to spend them on verbose wording
77
                nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "index must be int"));
78
            }
79
80
81
82
83
84

        case RT_BINARY_OP_ADD:
        case RT_BINARY_OP_INPLACE_ADD:
            if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
                // add 2 strings
                const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr);
85
86
87
                size_t lhs_len = strlen(lhs_str);
                size_t rhs_len = strlen(rhs_str);
                int alloc_len = lhs_len + rhs_len + 1;
88
                char *val = m_new(char, alloc_len);
89
90
91
                memcpy(val, lhs_str, lhs_len);
                memcpy(val + lhs_len, rhs_str, rhs_len);
                val[lhs_len + rhs_len] = '\0';
92
                return mp_obj_new_str(qstr_from_str_take(val, alloc_len));
93
94
            }
            break;
95
96
97
98
99
100
101
102
103
        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)) {
                const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr);
                /* FIXME \0 in strs */
                return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (strstr(lhs_str, rhs_str) == NULL));
            }
            break;
104
105
106
107
108
109
110
111
112
113
114
115
        case RT_BINARY_OP_MULTIPLY:
        {
            if (!MP_OBJ_IS_SMALL_INT(rhs_in)) {
                return NULL;
            }
            int n = MP_OBJ_SMALL_INT_VALUE(rhs_in);
            size_t len = strlen(lhs_str);
            char *s = m_new(char, len * n + 1);
            s[len * n] = 0;
            mp_seq_multiply(lhs_str, sizeof(*lhs_str), len, n, s);
            return MP_OBJ_NEW_QSTR(qstr_from_str_take(s, len * n + 1));
        }
116
117
118
119
120
    }

    return MP_OBJ_NULL; // op not supported
}

xyb's avatar
xyb committed
121
122
123
124
static mp_obj_t str_getiter(mp_obj_t o_in) {
    return mp_obj_new_str_iterator(o_in, 0);
}

125
126
127
128
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;

129
130
131
132
133
    // get separation string
    const char *sep_str = qstr_str(self->qstr);
    size_t sep_len = strlen(sep_str);

    // process args
134
135
136
137
138
139
140
141
142
    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;
    }
143
144
145

    // count required length
    int required_len = 0;
146
147
148
149
    for (int i = 0; i < seq_len; i++) {
        if (!MP_OBJ_IS_TYPE(seq_items[i], &str_type)) {
            goto bad_arg;
        }
150
151
152
        if (i > 0) {
            required_len += sep_len;
        }
153
154
155
156
157
        required_len += strlen(qstr_str(mp_obj_str_get(seq_items[i])));
    }

    // make joined string
    char *joined_str = m_new(char, required_len + 1);
158
    char *s_dest = joined_str;
159
160
    for (int i = 0; i < seq_len; i++) {
        if (i > 0) {
161
162
            memcpy(s_dest, sep_str, sep_len);
            s_dest += sep_len;
163
        }
164
165
166
167
        const char *s2 = qstr_str(mp_obj_str_get(seq_items[i]));
        size_t s2_len = strlen(s2);
        memcpy(s_dest, s2, s2_len);
        s_dest += s2_len;
168
    }
169
170
171
    *s_dest = '\0';

    // return joined string
172
    return mp_obj_new_str(qstr_from_str_take(joined_str, required_len + 1));
173
174

bad_arg:
175
    nlr_jump(mp_obj_new_exception_msg(MP_QSTR_TypeError, "?str.join expecting a list of str's"));
176
177
}

xbe's avatar
xbe committed
178
179
180
181
182
183
184
185
186
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;
}

187
static mp_obj_t str_find(uint n_args, const mp_obj_t *args) {
188
    assert(2 <= n_args && n_args <= 4);
189
190
    const char* haystack = qstr_str(mp_obj_str_get(args[0]));
    const char* needle = qstr_str(mp_obj_str_get(args[1]));
191

192
193
    size_t haystack_len = strlen(haystack);
    size_t needle_len = strlen(needle);
194
195
196
197
198
199
200
201
202
203
204
205

    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);
206
207
208
209
210
211
    if (p == NULL) {
        // not found
        return MP_OBJ_NEW_SMALL_INT(-1);
    } else {
        // found
        machine_int_t pos = p - haystack;
212
213
214
        if (pos + needle_len > end) {
            pos = -1;
        }
215
        return MP_OBJ_NEW_SMALL_INT(pos);
216
217
218
    }
}

219
mp_obj_t str_strip(uint n_args, const mp_obj_t *args) {
xbe's avatar
xbe committed
220
221
222
223
224
225
226
227
    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 {
228
        chars_to_del = qstr_str(mp_obj_str_get(args[1]));
xbe's avatar
xbe committed
229
230
231
    }

    const size_t chars_to_del_len = strlen(chars_to_del);
232
    const char *orig_str = qstr_str(mp_obj_str_get(args[0]));
xbe's avatar
xbe committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
    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) {
        //string is all whitespace, return '\0'
        char *empty = m_new(char, 1);
        empty[0] = '\0';
        return mp_obj_new_str(qstr_from_str_take(empty, 1));
    }

    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;
    //+1 to accomodate '\0'
    char *stripped_str = m_new(char, stripped_len + 1);
260
    memcpy(stripped_str, orig_str + first_good_char_pos, stripped_len);
xbe's avatar
xbe committed
261
262
263
264
    stripped_str[stripped_len] = '\0';
    return mp_obj_new_str(qstr_from_str_take(stripped_str, stripped_len + 1));
}

265
mp_obj_t str_format(uint n_args, const mp_obj_t *args) {
266
267
268
269
270
271
272
273
274
275
276
    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, '{');
277
278
            } else {
                while (*str != '}') str++;
279
                if (arg_i >= n_args) {
280
                    nlr_jump(mp_obj_new_exception_msg(MP_QSTR_IndexError, "tuple index out of range"));
281
                }
282
                // TODO: may be PRINT_REPR depending on formatting code
283
                mp_obj_print_helper((void (*)(void*, const char*, ...))vstr_printf, vstr, args[arg_i], PRINT_STR);
284
285
286
287
288
289
290
                arg_i++;
            }
        } else {
            vstr_add_char(vstr, *str);
        }
    }

291
    return mp_obj_new_str(qstr_from_str_take(vstr->buf, vstr->alloc));
292
293
}

294
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_find_obj, 2, 4, str_find);
295
static MP_DEFINE_CONST_FUN_OBJ_2(str_join_obj, str_join);
xbe's avatar
xbe committed
296
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(str_strip_obj, 1, 2, str_strip);
297
298
static MP_DEFINE_CONST_FUN_OBJ_VAR(str_format_obj, 1, str_format);

ian-v's avatar
ian-v committed
299
static const mp_method_t str_type_methods[] = {
300
    { "find", &str_find_obj },
ian-v's avatar
ian-v committed
301
    { "join", &str_join_obj },
xbe's avatar
xbe committed
302
    { "strip", &str_strip_obj },
ian-v's avatar
ian-v committed
303
304
305
    { "format", &str_format_obj },
    { NULL, NULL }, // end-of-list sentinel
};
306

307
308
309
const mp_obj_type_t str_type = {
    { &mp_const_type },
    "str",
310
311
312
    .print = str_print,
    .binary_op = str_binary_op,
    .getiter = str_getiter,
ian-v's avatar
ian-v committed
313
    .methods = str_type_methods,
314
315
316
317
318
319
320
321
322
323
};

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) {
324
325
326
327
328
329
330
331
332
    if (MP_OBJ_IS_QSTR(self_in)) {
        return MP_OBJ_QSTR_VALUE(self_in);
    }
    if (MP_OBJ_IS_TYPE(self_in, &str_type)) {
        mp_obj_str_t *self = self_in;
        return self->qstr;
    }
    nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError, "Can't convert '%s' object to str implicitly",
             mp_obj_get_type_str(self_in)));
333
}
xyb's avatar
xyb committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358

/******************************************************************************/
/* 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)) {
        mp_obj_t o_out = mp_obj_new_str(qstr_from_strn_copy(str + self->cur, 1));
        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",
359
    .iternext = str_it_iternext,
xyb's avatar
xyb committed
360
361
362
363
364
365
366
367
368
};

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;
}