|
| 1 | +#include <string.h> |
| 2 | + |
| 3 | +#include "misc.h" |
| 4 | +#include "mpconfig.h" |
| 5 | +#include "nlr.h" |
| 6 | +#include "qstr.h" |
| 7 | +#include "obj.h" |
| 8 | +#include "objtuple.h" |
| 9 | + |
| 10 | +typedef struct _mp_obj_namedtuple_type_t { |
| 11 | + mp_obj_type_t base; |
| 12 | + const char *fields; |
| 13 | +} mp_obj_namedtuple_type_t; |
| 14 | + |
| 15 | +typedef struct _mp_obj_namedtuple_t { |
| 16 | + mp_obj_tuple_t tuple; |
| 17 | +} mp_obj_namedtuple_t; |
| 18 | + |
| 19 | +static inline bool is_end_tok(char c) { |
| 20 | + return c == ' ' || c == ','; |
| 21 | +} |
| 22 | + |
| 23 | +static inline const char *skip_to_next(const char *p) { |
| 24 | + while (!is_end_tok(*p)) { |
| 25 | + if (*p == 0) { |
| 26 | + return NULL; |
| 27 | + } |
| 28 | + p++; |
| 29 | + } |
| 30 | + while (is_end_tok(*p)) { |
| 31 | + if (*p == 0) { |
| 32 | + return NULL; |
| 33 | + } |
| 34 | + p++; |
| 35 | + } |
| 36 | + return p; |
| 37 | +} |
| 38 | + |
| 39 | +STATIC uint namedtuple_count_fields(const char *namedef) { |
| 40 | + uint cnt = 0; |
| 41 | + while (*namedef != 0) { |
| 42 | + cnt++; |
| 43 | + while (!is_end_tok(*namedef) && *namedef != 0) { |
| 44 | + namedef++; |
| 45 | + } |
| 46 | + while (is_end_tok(*namedef) && *namedef != 0) { |
| 47 | + namedef++; |
| 48 | + } |
| 49 | + } |
| 50 | + return cnt; |
| 51 | +} |
| 52 | + |
| 53 | +STATIC int namedtuple_find_field(const char *name, const char *namedef) { |
| 54 | + int id = 0; |
| 55 | + int len = strlen(name); |
| 56 | + while (namedef) { |
| 57 | + if (memcmp(name, namedef, len) == 0) { |
| 58 | + namedef += len; |
| 59 | + if (*namedef == 0 || is_end_tok(*namedef)) { |
| 60 | + return id; |
| 61 | + } |
| 62 | + } |
| 63 | + namedef = skip_to_next(namedef); |
| 64 | + id++; |
| 65 | + } |
| 66 | + |
| 67 | + return -1; |
| 68 | +} |
| 69 | + |
| 70 | +STATIC void namedtuple_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t o_in, mp_print_kind_t kind) { |
| 71 | + mp_obj_namedtuple_t *o = o_in; |
| 72 | + print(env, "%s(", qstr_str(o->tuple.base.type->name)); |
| 73 | + const char *fields = ((mp_obj_namedtuple_type_t*)o->tuple.base.type)->fields; |
| 74 | + |
| 75 | + for (int i = 0; i < o->tuple.len; i++) { |
| 76 | + if (i > 0) { |
| 77 | + print(env, ", "); |
| 78 | + } |
| 79 | + const char *next = fields; |
| 80 | + |
| 81 | + while (!is_end_tok(*next) && *next != 0) { |
| 82 | + next++; |
| 83 | + } |
| 84 | + print(env, "%.*s=", next - fields, fields); |
| 85 | + mp_obj_print_helper(print, env, o->tuple.items[i], PRINT_REPR); |
| 86 | + while (is_end_tok(*next) && *next != 0) { |
| 87 | + next++; |
| 88 | + } |
| 89 | + fields = next; |
| 90 | + } |
| 91 | + print(env, ")"); |
| 92 | +} |
| 93 | + |
| 94 | +STATIC void namedtuple_load_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) { |
| 95 | + mp_obj_namedtuple_t *self = self_in; |
| 96 | + const char *fields = ((mp_obj_namedtuple_type_t*)self->tuple.base.type)->fields; |
| 97 | + int id = namedtuple_find_field(qstr_str(attr), fields); |
| 98 | + if (id < 0) { |
| 99 | + return; |
| 100 | + } |
| 101 | + dest[0] = self->tuple.items[id]; |
| 102 | +} |
| 103 | + |
| 104 | +STATIC bool namedtuple_store_attr(mp_obj_t self_in, qstr attr, mp_obj_t value) { |
| 105 | + nlr_jump(mp_obj_new_exception_msg(&mp_type_AttributeError, "can't set attribute")); |
| 106 | +} |
| 107 | + |
| 108 | +STATIC mp_obj_t namedtuple_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_obj_t *args) { |
| 109 | + mp_obj_namedtuple_type_t *type = type_in; |
| 110 | + uint num_fields = namedtuple_count_fields(type->fields); |
| 111 | + if (n_args != num_fields) { |
| 112 | + // Counts include implicit "self" |
| 113 | + nlr_jump(mp_obj_new_exception_msg_varg(&mp_type_TypeError, |
| 114 | + "__new__() takes %d positional arguments but %d were given", |
| 115 | + num_fields + 1, n_args + 1)); |
| 116 | + } |
| 117 | + mp_obj_tuple_t *tuple = mp_obj_new_tuple(n_args, args); |
| 118 | + tuple->base.type = type_in; |
| 119 | + return tuple; |
| 120 | +} |
| 121 | + |
| 122 | +STATIC const mp_obj_tuple_t namedtuple_base_tuple = {{&tuple_type}, 1, {(mp_obj_t)&tuple_type}}; |
| 123 | + |
| 124 | +mp_obj_t mp_obj_new_namedtuple_type(qstr name, const char *fields) { |
| 125 | + mp_obj_namedtuple_type_t *o = m_new0(mp_obj_namedtuple_type_t, 1); |
| 126 | + o->base.base.type = &mp_type_type; |
| 127 | + o->base.name = name; |
| 128 | + o->base.print = namedtuple_print; |
| 129 | + o->base.make_new = namedtuple_make_new; |
| 130 | +// o->base.unary_op = ; |
| 131 | +// o->base.binary_op = ; |
| 132 | + o->base.load_attr = namedtuple_load_attr; |
| 133 | + o->base.store_attr = namedtuple_store_attr; |
| 134 | + o->base.bases_tuple = (mp_obj_t)&namedtuple_base_tuple; |
| 135 | + o->fields = fields; |
| 136 | + return o; |
| 137 | +} |
| 138 | + |
| 139 | +STATIC mp_obj_t new_namedtuple_type(mp_obj_t name_in, mp_obj_t fields_in) { |
| 140 | + qstr name = MP_OBJ_QSTR_VALUE(name_in); |
| 141 | + const char *fields = mp_obj_str_get_str(fields_in); |
| 142 | + return mp_obj_new_namedtuple_type(name, fields); |
| 143 | +} |
| 144 | +MP_DEFINE_CONST_FUN_OBJ_2(mp_namedtuple_obj, new_namedtuple_type); |
0 commit comments