Skip to content

Commit 0fcbaa4

Browse files
chipacachipaca
committed
implemented dict.pop
1 parent cd08873 commit 0fcbaa4

4 files changed

Lines changed: 67 additions & 13 deletions

File tree

‎py/map.c‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,13 @@ static void mp_map_rehash (mp_map_t *map) {
5858
map->table = m_new0(mp_map_elem_t, map->alloc);
5959
for (int i = 0; i < old_alloc; i++) {
6060
if (old_table[i].key != NULL) {
61-
mp_map_lookup_helper(map, old_table[i].key, true)->value = old_table[i].value;
61+
mp_map_lookup_helper(map, old_table[i].key, true, false)->value = old_table[i].value;
6262
}
6363
}
6464
m_del(mp_map_elem_t, old_table, old_alloc);
6565
}
6666

67-
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found) {
67+
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) {
6868
bool is_map_mp_obj = (map->kind == MP_MAP_OBJ);
6969
machine_uint_t hash;
7070
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
105105
elem->key = index;
106106
}
107107
*/
108+
if (remove_if_found) {
109+
map->used--;
110+
/* this leaks this memory (but see dict_get_helper) */
111+
mp_map_elem_t *retval = m_new(mp_map_elem_t, 1);
112+
retval->key = elem->key;
113+
retval->value = elem->value;
114+
elem->key = NULL;
115+
return retval;
116+
}
108117
return elem;
109118
} else {
110119
// 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
115124

116125
mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found) {
117126
mp_obj_t o = (mp_obj_t)(machine_uint_t)index;
118-
return mp_map_lookup_helper(map, o, add_if_not_found);
127+
return mp_map_lookup_helper(map, o, add_if_not_found, false);
119128
}
120129

121130
/******************************************************************************/

‎py/map.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ typedef struct _mp_set_t {
2626
int get_doubling_prime_greater_or_equal_to(int x);
2727
void mp_map_init(mp_map_t *map, mp_map_kind_t kind, int n);
2828
mp_map_t *mp_map_new(mp_map_kind_t kind, int n);
29-
mp_map_elem_t* mp_map_lookup_helper(mp_map_t *map, mp_obj_t index, bool add_if_not_found);
29+
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);
3030
mp_map_elem_t* mp_qstr_map_lookup(mp_map_t *map, qstr index, bool add_if_not_found);
3131
void mp_map_clear(mp_map_t *map);
3232

‎py/objdict.c‎

Lines changed: 42 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
5050
case RT_BINARY_OP_SUBSCR:
5151
{
5252
// dict load
53-
mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false);
53+
mp_map_elem_t *elem = mp_map_lookup_helper(&o->map, rhs_in, false, false);
5454
if (elem == NULL) {
5555
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
5656
} else {
@@ -140,20 +140,52 @@ 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) {
146+
if (deflt == NULL) {
147+
if (pop) {
148+
nlr_jump(mp_obj_new_exception_msg(MP_QSTR_KeyError, "<value>"));
149+
} else {
150+
return mp_const_none;
151+
}
152+
} else {
153+
return deflt;
154+
}
155+
} else {
156+
mp_obj_t value = elem->value;
157+
if (pop) {
158+
/* catch the leak (from mp_map_lookup_helper) */
159+
m_free(elem, 2 * sizeof(mp_obj_t));
160+
}
161+
return value;
162+
}
163+
164+
}
165+
143166
static mp_obj_t dict_get(int n_args, const mp_obj_t *args) {
144167
assert(2 <= n_args && n_args <= 3);
145168
assert(MP_OBJ_IS_TYPE(args[0], &dict_type));
146169

147-
mp_map_elem_t *elem = mp_map_lookup_helper(&((mp_obj_dict_t *)args[0])->map,
148-
args[1], false);
149-
if (elem == NULL) {
150-
return n_args >= 3 ? args[2] : mp_const_none;
151-
} else {
152-
return elem->value;
153-
}
170+
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
171+
args[1],
172+
n_args == 3 ? args[2] : NULL,
173+
false);
154174
}
155175
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_get_obj, 2, 3, dict_get);
156176

177+
static mp_obj_t dict_pop(int n_args, const mp_obj_t *args) {
178+
assert(2 <= n_args && n_args <= 3);
179+
assert(MP_OBJ_IS_TYPE(args[0], &dict_type));
180+
181+
return dict_get_helper(&((mp_obj_dict_t *)args[0])->map,
182+
args[1],
183+
n_args == 3 ? args[2] : NULL,
184+
true);
185+
}
186+
static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(dict_pop_obj, 2, 3, dict_pop);
187+
188+
157189
/******************************************************************************/
158190
/* dict constructors & etc */
159191

@@ -168,6 +200,7 @@ const mp_obj_type_t dict_type = {
168200
{ "clear", &dict_clear_obj },
169201
{ "copy", &dict_copy_obj },
170202
{ "get", &dict_get_obj },
203+
{ "pop", &dict_pop_obj },
171204
{ NULL, NULL }, // end-of-list sentinel
172205
},
173206
};
@@ -186,6 +219,6 @@ uint mp_obj_dict_len(mp_obj_t self_in) {
186219
mp_obj_t mp_obj_dict_store(mp_obj_t self_in, mp_obj_t key, mp_obj_t value) {
187220
assert(MP_OBJ_IS_TYPE(self_in, &dict_type));
188221
mp_obj_dict_t *self = self_in;
189-
mp_map_lookup_helper(&self->map, key, true)->value = value;
222+
mp_map_lookup_helper(&self->map, key, true, false)->value = value;
190223
return self_in;
191224
}

‎tests/basics/tests/dict_pop.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
d = {1: 2, 3: 4}
2+
print(d.pop(3), d)
3+
print(d)
4+
print(d.pop(1, 42), d)
5+
print(d.pop(1, 42), d)
6+
print(d.pop(1, None), d)
7+
try:
8+
print(d.pop(1), "!!!",)
9+
except KeyError:
10+
print("Raised KeyError")
11+
else:
12+
print("Did not rise KeyError!")

0 commit comments

Comments
 (0)