Add built-in dict.pop

Author: youknowone

tests/snippets/dict.py

@@ -166,3 +166,11 @@ def __eq__(self, other):
 
 y.update(y)
 assert y == {'a': 2, 'b': 12, 'c': 19, 'd': -1}  # hasn't changed
+
+x = {1: 'a', '1': None}
+assert x.pop(1) == 'a'
+assert x.pop('1') is None
+assert x == {}
+
+with assertRaises(KeyError):
+    x.pop("not here")

vm/src/dictdatatype.rs

@@ -87,14 +87,18 @@ impl<T: Clone> Dict<T> {
         }
     }
 
+    fn unchecked_get(&self, index: usize) -> T {
+        if let Some(entry) = &self.entries[index] {
+            entry.value.clone()
+        } else {
+            panic!("Lookup returned invalid index into entries!");
+        }
+    }
+
     /// Retrieve a key
     pub fn get(&self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<Option<T>> {
         if let LookupResult::Existing(index) = self.lookup(vm, key)? {
-            if let Some(entry) = &self.entries[index] {
-                Ok(Some(entry.value.clone()))
-            } else {
-                panic!("Lookup returned invalid index into entries!");
-            }
+            Ok(Some(self.unchecked_get(index)))
         } else {
             Ok(None)
         }
@@ -190,6 +194,19 @@ impl<T: Clone> Dict<T> {
             // warn!("Perturb value: {}", i);
         }
     }
+
+    /// Retrieve and delete a key
+    pub fn pop(&mut self, vm: &VirtualMachine, key: &PyObjectRef) -> PyResult<T> {
+        if let LookupResult::Existing(index) = self.lookup(vm, key)? {
+            let value = self.unchecked_get(index);
+            self.entries[index] = None;
+            self.size -= 1;
+            Ok(value)
+        } else {
+            let key_repr = vm.to_pystr(key)?;
+            Err(vm.new_key_error(format!("Key not found: {}", key_repr)))
+        }
+    }
 }
 
 enum LookupResult {

vm/src/obj/objdict.rs

@@ -228,6 +228,10 @@ impl PyDictRef {
         PyDictRef::merge(&self.entries, dict_obj, kwargs, vm)
     }
 
+    fn pop(self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+        self.entries.borrow_mut().pop(vm, &key)
+    }
+
     /// Take a python dictionary and convert it to attributes.
     pub fn to_attributes(self) -> PyAttributes {
         let mut attrs = PyAttributes::new();
@@ -458,6 +462,7 @@ pub fn init(context: &PyContext) {
         "get" => context.new_rustfunc(PyDictRef::get),
         "copy" => context.new_rustfunc(PyDictRef::copy),
         "update" => context.new_rustfunc(PyDictRef::update),
+        "pop" => context.new_rustfunc(PyDictRef::pop),
     });
 
     PyDictKeys::extend_class(context, &context.dictkeys_type);