Skip to content

Commit be8fe5b

Browse files
chipacachipaca
committed
Added dict.setdefault
1 parent f77dce8 commit be8fe5b

3 files changed

Lines changed: 40 additions & 10 deletions

File tree

‎py/map.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_n
112112
retval->key = elem->key;
113113
retval->value = elem->value;
114114
elem->key = NULL;
115+
elem->value = NULL;
115116
return retval;
116117
}
117118
return elem;

‎py/objdict.c‎

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -140,27 +140,30 @@ static mp_obj_t dict_copy(mp_obj_t self_in) {
140140
}
141141
static MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy);
142142

143-
static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, bool pop) {
144-
mp_map_elem_t *elem = mp_map_lookup_helper(self, key, false, pop);
145-
if (elem == NULL) {
143+
static mp_obj_t dict_get_helper(mp_map_t *self, mp_obj_t key, mp_obj_t deflt, bool pop, bool set) {
144+
mp_map_elem_t *elem = mp_map_lookup_helper(self, key, set, pop);
145+
mp_obj_t value;
146+
if (elem == NULL || elem->value == NULL) {
146147
if (deflt == NULL) {
147148
if (pop) {
148149
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
149150
} else {
150-
return mp_const_none;
151+
value = mp_const_none;
151152
}
152153
} else {
153-
return deflt;
154+
value = deflt;
154155
}
155156
} else {
156-
mp_obj_t value = elem->value;
157+
value = elem->value;
157158
if (pop) {
158159
/* catch the leak (from mp_map_lookup_helper) */
159160
m_free(elem, 2 * sizeof(mp_obj_t));
160161
}
161-
return value;
162162
}
163-
163+
if (set) {
164+
elem->value = value;
165+
}
166+
return value;
164167
}
165168

166169
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) {
170173
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
171174
args[1],
172175
n_args == 3 ? args[2] : NULL,
173-
false);
176+
false, false);
174177
}
175178
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);
176179

@@ -181,11 +184,22 @@ static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) {
181184
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
182185
args[1],
183186
n_args == 3 ? args[2] : NULL,
184-
true);
187+
true, false);
185188
}
186189
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);
187190

188191

192+
static mp_obj_t dict_setdefault(int n_args, const mp_obj_t *args) {
193+
assert(2 <= n_args && n_args <= 3);
194+
assert(MP_OBJ_IS_TYPE(args[0], &dict_type));
195+
196+
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
197+
args[1],
198+
n_args == 3 ? args[2] : NULL,
199+
false, true);
200+
}
201+
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_setdefault_obj, 2, 3, dict_setdefault);
202+
189203

190204
static mp_obj_t dict_popitem(mp_obj_t self_in) {
191205
assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
@@ -199,6 +213,7 @@ static mp_obj_t dict_popitem(mp_obj_t self_in) {
199213
self->map.used--;
200214
mp_obj_t items[] = {next->key, next->value};
201215
next->key = NULL;
216+
next->value = NULL;
202217
mp_obj_t tuple = mp_obj_new_tuple(2, items);
203218

204219
return tuple;
@@ -222,6 +237,7 @@ const mp_obj_type_t dict_type = {
222237
{ "get", &dict_get_obj },
223238
{ "pop", &dict_pop_obj },
224239
{ "popitem", &dict_popitem_obj },
240+
{ "setdefault", &dict_setdefault_obj },
225241
{ NULL, NULL }, // end-of-list sentinel
226242
},
227243
};

‎tests/basics/tests/dict_setdefault.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
d = {}
2+
print(d.setdefault(1))
3+
print(d.setdefault(1))
4+
print(d.setdefault(5, 42))
5+
print(d.setdefault(5, 1))
6+
print(d[1])
7+
print(d[5])
8+
d.pop(5)
9+
print(d.setdefault(5, 1))
10+
print(d[1])
11+
print(d[5])
12+
13+

0 commit comments

Comments
 (0)