implemented dict.pop

4 files changed, 67 insertions, 13 deletions

diff --git a/py/map.c b/py/map.c
index dff0bdfea0..8554150adf 100644
--- a/py/map.c
+++ b/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);
 }
 
 /******************************************************************************/
diff --git a/py/map.h b/py/map.h
index b4592c270a..f0a57758c5 100644
--- a/py/map.h
+++ b/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);
 
diff --git a/py/objdict.c b/py/objdict.c
index fd5e37dfe0..935829ddb5 100644
--- a/py/objdict.c
+++ b/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;
 }
diff --git a/tests/basics/tests/dict_pop.py b/tests/basics/tests/dict_pop.py
new file mode 100644
index 0000000000..602560ce9d
--- /dev/null
+++ b/tests/basics/tests/dict_pop.py
@@ -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!")