Add dict.setdefault.

Author: cthulahoops

tests/snippets/dict.py

@@ -180,3 +180,11 @@ def __eq__(self, other):
 with assertRaises(KeyError):
     x.popitem()
 assert x == {}
+
+x = {'a': 4}
+assert 4 == x.setdefault('a', 0)
+assert x['a'] == 4
+assert 0 == x.setdefault('b', 0)
+assert x['b'] == 0
+assert None == x.setdefault('c')
+assert x['c'] is None

vm/src/obj/objdict.rs

@@ -213,6 +213,26 @@ impl PyDictRef {
         }
     }
 
+    fn setdefault(
+        self,
+        key: PyObjectRef,
+        default: OptionalArg<PyObjectRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        let mut entries = self.entries.borrow_mut();
+        match entries.get(vm, &key)? {
+            Some(value) => Ok(value),
+            None => {
+                let set_value = match default {
+                    OptionalArg::Present(value) => value,
+                    OptionalArg::Missing => vm.ctx.none(),
+                };
+                entries.insert(vm, &key, set_value.clone())?;
+                Ok(set_value)
+            }
+        }
+    }
+
     fn copy(self, _vm: &VirtualMachine) -> PyDict {
         PyDict {
             entries: self.entries.clone(),
@@ -473,6 +493,7 @@ pub fn init(context: &PyContext) {
         "items" => context.new_rustfunc(PyDictRef::items),
         "keys" => context.new_rustfunc(PyDictRef::keys),
         "get" => context.new_rustfunc(PyDictRef::get),
+        "setdefault" => context.new_rustfunc(PyDictRef::setdefault),
         "copy" => context.new_rustfunc(PyDictRef::copy),
         "update" => context.new_rustfunc(PyDictRef::update),
         "pop" => context.new_rustfunc(PyDictRef::pop),