diff options
Diffstat (limited to 'py/objdict.c')
| -rw-r--r-- | py/objdict.c | 233 |
1 files changed, 214 insertions, 19 deletions
diff --git a/py/objdict.c b/py/objdict.c index 66a442fbaa..c36d4ccb03 100644 --- a/py/objdict.c +++ b/py/objdict.c @@ -17,20 +17,23 @@ typedef struct _mp_obj_dict_t { mp_map_t map; } mp_obj_dict_t; +static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur); +static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in); + static void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in) { mp_obj_dict_t *self = self_in; bool first = true; print(env, "{"); - for (int i = 0; i < self->map.alloc; i++) { - if (self->map.table[i].key != NULL) { - if (!first) { - print(env, ", "); - } - first = false; - mp_obj_print_helper(print, env, self->map.table[i].key); - print(env, ": "); - mp_obj_print_helper(print, env, self->map.table[i].value); + mp_obj_t *dict_iter = mp_obj_new_dict_iterator(self, 0); + mp_map_elem_t *next = NULL; + while ((next = dict_it_iternext_elem(dict_iter)) != NULL) { + if (!first) { + print(env, ", "); } + first = false; + mp_obj_print_helper(print, env, next->key); + print(env, ": "); + mp_obj_print_helper(print, env, next->value); } print(env, "}"); } @@ -47,7 +50,7 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { case RT_BINARY_OP_SUBSCR: { // dict load - mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false); + mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false, false); if (elem == NULL) { nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>")); } else { @@ -60,12 +63,211 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { } } + +/******************************************************************************/ +/* dict iterator */ + +typedef struct _mp_obj_dict_it_t { + mp_obj_base_t base; + mp_obj_dict_t *dict; + machine_uint_t cur; +} mp_obj_dict_it_t; + +static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) { + mp_obj_dict_it_t *self = self_in; + machine_uint_t max = self->dict->map.alloc; + mp_map_elem_t *table = self->dict->map.table; + + for (int i = self->cur; i < max; i++) { + if (table[i].key != NULL) { + self->cur = i + 1; + return &(table[i]); + } + } + + return NULL; +} + +mp_obj_t dict_it_iternext(mp_obj_t self_in) { + mp_map_elem_t *next = dict_it_iternext_elem(self_in); + + if (next != NULL) { + return next->key; + } else { + return mp_const_stop_iteration; + } +} + +static const mp_obj_type_t dict_it_type = { + { &mp_const_type }, + "dict_iterator", + .iternext = dict_it_iternext, +}; + +static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur) { + mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t); + o->base.type = &dict_it_type; + o->dict = dict; + o->cur = cur; + return o; +} + +static mp_obj_t dict_getiter(mp_obj_t o_in) { + return mp_obj_new_dict_iterator(o_in, 0); +} + +/******************************************************************************/ +/* dict methods */ + +static mp_obj_t dict_clear(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + + mp_map_clear(&self->map); + + return mp_const_none; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_clear_obj, dict_clear); + +static mp_obj_t dict_copy(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + mp_obj_dict_t *other = mp_obj_new_dict(self->map.alloc); + other->map.used = self->map.used; + memcpy(other->map.table, self->map.table, self->map.alloc * sizeof(mp_map_elem_t)); + return other; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy); + +static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, bool pop, bool set) { + mp_map_elem_t *elem = mp_map_lookup_helper(self, key, set, pop); + mp_obj_t value; + if (elem == NULL || elem->value == NULL) { + if (deflt == NULL) { + if (pop) { + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>")); + } else { + value = mp_const_none; + } + } else { + value = deflt; + } + } else { + value = elem->value; + if (pop) { + /* catch the leak (from mp_map_lookup_helper) */ + m_free(elem, sizeof(mp_map_elem_t)); + } + } + if (set) { + elem->value = value; + } + return value; +} + +static mp_obj_t dict_get(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + false, false); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get); + +static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + true, false); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop); + + +static mp_obj_t dict_setdefault(int n_args, const mp_obj_t *args) { + assert(2 <= n_args && n_args <= 3); + assert(MP_OBJ_IS_TYPE(args[0], &dict_type)); + + return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, + args[1], + n_args == 3 ? args[2] : NULL, + false, true); +} +static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault); + + +static mp_obj_t dict_popitem(mp_obj_t self_in) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + if (self->map.used == 0) { + nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "popitem(): dictionary is empty")); + } + mp_obj_dict_it_t *iter = mp_obj_new_dict_iterator(self, 0); + + mp_map_elem_t *next = dict_it_iternext_elem(iter); + self->map.used--; + mp_obj_t items[] = {next->key, next->value}; + next->key = NULL; + next->value = NULL; + mp_obj_t tuple = mp_obj_new_tuple(2, items); + + return tuple; +} +static MP_DEFINE_CONST_FUN_OBJ_1(dict_popitem_obj, dict_popitem); + +static mp_obj_t dict_update(mp_obj_t self_in, mp_obj_t iterable) { + assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); + mp_obj_dict_t *self = self_in; + /* TODO: check for the "keys" method */ + mp_obj_t iter = rt_getiter(iterable); + mp_obj_t next = NULL; + while ((next = rt_iternext(iter)) != mp_const_stop_iteration) { + mp_obj_t inneriter = rt_getiter(next); + mp_obj_t key = rt_iternext(inneriter); + mp_obj_t value = rt_iternext(inneriter); + mp_obj_t stop = rt_iternext(inneriter); + if (key == mp_const_stop_iteration + || value == mp_const_stop_iteration + || stop != mp_const_stop_iteration) { + nlr_jump(mp_obj_new_exception_msg( + MP_QSTR_ValueError, + "dictionary update sequence has the wrong length")); + } else { + mp_map_lookup_helper(&self->map, key, true, false)->value = value; + } + } + + return mp_const_none; +} +static MP_DEFINE_CONST_FUN_OBJ_2(dict_update_obj, dict_update); + + +/******************************************************************************/ +/* dict constructors & etc */ + +static const mp_method_t dict_type_methods[] = { + { "clear", &dict_clear_obj }, + { "copy", &dict_copy_obj }, + { "get", &dict_get_obj }, + { "pop", &dict_pop_obj }, + { "popitem", &dict_popitem_obj }, + { "setdefault", &dict_setdefault_obj }, + { "update", &dict_update_obj }, + { NULL, NULL }, // end-of-list sentinel +}; + const mp_obj_type_t dict_type = { { &mp_const_type }, "dict", .print = dict_print, .make_new = dict_make_new, .binary_op = dict_binary_op, + .getiter = dict_getiter, + .methods = dict_type_methods, }; mp_obj_t mp_obj_new_dict(int n_args) { @@ -76,19 +278,12 @@ mp_obj_t mp_obj_new_dict(int n_args) { } uint mp_obj_dict_len(mp_obj_t self_in) { - mp_obj_dict_t *self = self_in; - uint len = 0; - for (int i = 0; i < self->map.alloc; i++) { - if (self->map.table[i].key != NULL) { - len += 1; - } - } - return len; + return ((mp_obj_dict_t *)self_in)->map.used; } mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value) { assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); mp_obj_dict_t *self = self_in; - mp_map_lookup_helper(&self->map, key, true)->value = value; + mp_map_lookup_helper(&self->map, key, true, false)->value = value; return self_in; } |