Skip to content
Snippets Groups Projects
Commit be8fe5be authored by John R. Lenton's avatar John R. Lenton
Browse files

Added dict.setdefault

parent f77dce8a
No related branches found
No related tags found
No related merge requests found
...@@ -112,6 +112,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n ...@@ -112,6 +112,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n
retval->key = elem->key; retval->key = elem->key;
retval->value = elem->value; retval->value = elem->value;
elem->key = NULL; elem->key = NULL;
elem->value = NULL;
return retval; return retval;
} }
return elem; return elem;
......
...@@ -140,27 +140,30 @@ static mp_obj_t dict_copy(mp_obj_t self_in) { ...@@ -140,27 +140,30 @@ 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_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) { 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, false, pop); mp_map_elem_t *elem = mp_map_lookup_helper(self, key, set, pop);
if (elem == NULL) { mp_obj_t value;
if (elem == NULL || elem->value == NULL) {
if (deflt == NULL) { if (deflt == NULL) {
if (pop) { if (pop) {
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>")); nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
} else { } else {
return mp_const_none; value = mp_const_none;
} }
} else { } else {
return deflt; value = deflt;
} }
} else { } else {
mp_obj_t value = elem->value; value = elem->value;
if (pop) { if (pop) {
/* catch the leak (from mp_map_lookup_helper) */ /* catch the leak (from mp_map_lookup_helper) */
m_free(elem, 2 * sizeof(mp_obj_t)); m_free(elem, 2 * sizeof(mp_obj_t));
} }
return value;
} }
if (set) {
elem->value = value;
}
return value;
} }
static mp_obj_t dict_get(int n_args, const mp_obj_t *args) { static mp_obj_t dict_get(int n_args, const mp_obj_t *args) {
...@@ -170,7 +173,7 @@ static mp_obj_t dict_get(int n_args, const mp_obj_t *args) { ...@@ -170,7 +173,7 @@ static mp_obj_t dict_get(int n_args, const mp_obj_t *args) {
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
args[1], args[1],
n_args == 3 ? args[2] : NULL, n_args == 3 ? args[2] : NULL,
false); false, false);
} }
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get); static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);
...@@ -181,11 +184,22 @@ static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) { ...@@ -181,11 +184,22 @@ static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) {
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
args[1], args[1],
n_args == 3 ? args[2] : NULL, n_args == 3 ? args[2] : NULL,
true); true, false);
} }
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop); 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) { static mp_obj_t dict_popitem(mp_obj_t self_in) {
assert(MP_OBJ_IS_TYPE(self_in, &dict_type)); assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
...@@ -199,6 +213,7 @@ static mp_obj_t dict_popitem(mp_obj_t self_in) { ...@@ -199,6 +213,7 @@ static mp_obj_t dict_popitem(mp_obj_t self_in) {
self->map.used--; self->map.used--;
mp_obj_t items[] = {next->key, next->value}; mp_obj_t items[] = {next->key, next->value};
next->key = NULL; next->key = NULL;
next->value = NULL;
mp_obj_t tuple = mp_obj_new_tuple(2, items); mp_obj_t tuple = mp_obj_new_tuple(2, items);
return tuple; return tuple;
...@@ -222,6 +237,7 @@ const mp_obj_type_t dict_type = { ...@@ -222,6 +237,7 @@ const mp_obj_type_t dict_type = {
{ "get", &dict_get_obj }, { "get", &dict_get_obj },
{ "pop", &dict_pop_obj }, { "pop", &dict_pop_obj },
{ "popitem", &dict_popitem_obj }, { "popitem", &dict_popitem_obj },
{ "setdefault", &dict_setdefault_obj },
{ NULL, NULL }, // end-of-list sentinel { NULL, NULL }, // end-of-list sentinel
}, },
}; };
......
d = {}
print(d.setdefault(1))
print(d.setdefault(1))
print(d.setdefault(5, 42))
print(d.setdefault(5, 1))
print(d[1])
print(d[5])
d.pop(5)
print(d.setdefault(5, 1))
print(d[1])
print(d[5])
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment