diff --git a/py/map.c b/py/map.c index 8554150adf..b88181989d 100644 --- a/py/map.c +++ b/py/map.c @@ -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->value = elem->value; elem->key = NULL; + elem->value = NULL; return retval; } return elem; diff --git a/py/objdict.c b/py/objdict.c index b4cf8f4aa8..753db7d808 100644 --- a/py/objdict.c +++ b/py/objdict.c @@ -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_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) { +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, "")); } else { - return mp_const_none; + value = mp_const_none; } } else { - return deflt; + value = deflt; } } else { - mp_obj_t value = elem->value; + value = elem->value; if (pop) { /* catch the leak (from mp_map_lookup_helper) */ 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) { @@ -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, args[1], n_args == 3 ? args[2] : NULL, - false); + false, false); } 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) { return dict_get_helper(&((mp_obj_dict_t *)args[0])->map, args[1], 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_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)); @@ -199,6 +213,7 @@ static mp_obj_t dict_popitem(mp_obj_t self_in) { 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; @@ -222,6 +237,7 @@ const mp_obj_type_t dict_type = { { "get", &dict_get_obj }, { "pop", &dict_pop_obj }, { "popitem", &dict_popitem_obj }, + { "setdefault", &dict_setdefault_obj }, { NULL, NULL }, // end-of-list sentinel }, }; diff --git a/tests/basics/tests/dict_setdefault.py b/tests/basics/tests/dict_setdefault.py new file mode 100644 index 0000000000..57d0ba4518 --- /dev/null +++ b/tests/basics/tests/dict_setdefault.py @@ -0,0 +1,13 @@ +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]) + +