implemented dict.pop
This commit is contained in:
parent
cd0887352d
commit
0fcbaa442f
15
py/map.c
15
py/map.c
|
@ -58,13 +58,13 @@ static void mp_map_rehash (mp_map_t *map) {
|
|||
map->table = m_new0(mp_map_elem_t, map->alloc);
|
||||
for (int i = 0; i < old_alloc; i++) {
|
||||
if (old_table[i].key != NULL) {
|
||||
mp_map_lookup_helper(map, old_table[i].key, true)->value = old_table[i].value;
|
||||
mp_map_lookup_helper(map, old_table[i].key, true, false)->value = old_table[i].value;
|
||||
}
|
||||
}
|
||||
m_del(mp_map_elem_t, old_table, old_alloc);
|
||||
}
|
||||
|
||||
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found) {
|
||||
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found, bool remove_if_found) {
|
||||
bool is_map_mp_obj = (map->kind == MP_MAP_OBJ);
|
||||
machine_uint_t hash;
|
||||
if (is_map_mp_obj) {
|
||||
|
@ -105,6 +105,15 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n
|
|||
elem->key = index;
|
||||
}
|
||||
*/
|
||||
if (remove_if_found) {
|
||||
map->used--;
|
||||
/* this leaks this memory (but see dict_get_helper) */
|
||||
mp_map_elem_t *retval = m_new(mp_map_elem_t, 1);
|
||||
retval->key = elem->key;
|
||||
retval->value = elem->value;
|
||||
elem->key = NULL;
|
||||
return retval;
|
||||
}
|
||||
return elem;
|
||||
} else {
|
||||
// not yet found, keep searching in this table
|
||||
|
@ -115,7 +124,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n
|
|||
|
||||
mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found) {
|
||||
mp_obj_t o = (mp_obj_t)(machine_uint_t)index;
|
||||
return mp_map_lookup_helper(map, o, add_if_not_found);
|
||||
return mp_map_lookup_helper(map, o, add_if_not_found, false);
|
||||
}
|
||||
|
||||
/******************************************************************************/
|
||||
|
|
2
py/map.h
2
py/map.h
|
@ -26,7 +26,7 @@ typedef struct _mp_set_t {
|
|||
int get_doubling_prime_greater_or_equal_to(int x);
|
||||
void mp_map_init(mp_map_t *map, mp_map_kind_t kind, int n);
|
||||
mp_map_t *mp_map_new(mp_map_kind_t kind, int n);
|
||||
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found);
|
||||
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found, bool remove_if_found);
|
||||
mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found);
|
||||
void mp_map_clear(mp_map_t *map);
|
||||
|
||||
|
|
51
py/objdict.c
51
py/objdict.c
|
@ -50,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 {
|
||||
|
@ -140,20 +140,52 @@ static mp_obj_t dict_copy(mp_obj_t self_in) {
|
|||
}
|
||||
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) {
|
||||
mp_map_elem_t *elem = mp_map_lookup_helper(self, key, false, pop);
|
||||
if (elem == NULL) {
|
||||
if (deflt == NULL) {
|
||||
if (pop) {
|
||||
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
|
||||
} else {
|
||||
return mp_const_none;
|
||||
}
|
||||
} else {
|
||||
return deflt;
|
||||
}
|
||||
} else {
|
||||
mp_obj_t value = elem->value;
|
||||
if (pop) {
|
||||
/* catch the leak (from mp_map_lookup_helper) */
|
||||
m_free(elem, 2 * sizeof(mp_obj_t));
|
||||
}
|
||||
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));
|
||||
|
||||
mp_map_elem_t *elem = mp_map_lookup_helper(&((mp_obj_dict_t *)args[0])->map,
|
||||
args[1], false);
|
||||
if (elem == NULL) {
|
||||
return n_args >= 3 ? args[2] : mp_const_none;
|
||||
} else {
|
||||
return elem->value;
|
||||
}
|
||||
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
|
||||
args[1],
|
||||
n_args == 3 ? args[2] : NULL,
|
||||
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);
|
||||
}
|
||||
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);
|
||||
|
||||
|
||||
/******************************************************************************/
|
||||
/* dict constructors & etc */
|
||||
|
||||
|
@ -168,6 +200,7 @@ const mp_obj_type_t dict_type = {
|
|||
{ "clear", &dict_clear_obj },
|
||||
{ "copy", &dict_copy_obj },
|
||||
{ "get", &dict_get_obj },
|
||||
{ "pop", &dict_pop_obj },
|
||||
{ NULL, NULL }, // end-of-list sentinel
|
||||
},
|
||||
};
|
||||
|
@ -186,6 +219,6 @@ uint mp_obj_dict_len(mp_obj_t self_in) {
|
|||
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;
|
||||
}
|
||||
|
|
|
@ -0,0 +1,12 @@
|
|||
d = {1: 2, 3: 4}
|
||||
print(d.pop(3), d)
|
||||
print(d)
|
||||
print(d.pop(1, 42), d)
|
||||
print(d.pop(1, 42), d)
|
||||
print(d.pop(1, None), d)
|
||||
try:
|
||||
print(d.pop(1), "!!!",)
|
||||
except KeyError:
|
||||
print("Raised KeyError")
|
||||
else:
|
||||
print("Did not rise KeyError!")
|
Loading…
Reference in New Issue