Merge pull request RustPython#608 from palaviv/dict-methods

Add dict.{keys,values,items}

Author: rmliddle

tests/snippets/dict.py

@@ -17,3 +17,20 @@ def dict_eq(d1, d2):
 
 a.clear()
 assert len(a) == 0
+
+a = {'a': 5, 'b': 6}
+res = set()
+for value in a.values():
+	res.add(value)
+assert res == set([5,6])
+
+count = 0
+for (key, value) in a.items():
+	assert a[key] == value
+	count += 1
+assert count == len(a)
+
+res = set()
+for key in a.keys():
+	res.add(key)
+assert res == set(['a','b'])

vm/src/obj/objdict.rs

@@ -243,8 +243,8 @@ fn dict_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(dict, Some(vm.ctx.dict_type()))]);
 
     let keys = get_elements(dict)
-        .keys()
-        .map(|k| vm.ctx.new_str(k.to_string()))
+        .values()
+        .map(|(k, _v)| k.clone())
         .collect();
     let key_list = vm.ctx.new_list(keys);
 
@@ -259,6 +259,46 @@ fn dict_iter(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(iter_obj)
 }
 
+fn dict_values(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(dict, Some(vm.ctx.dict_type()))]);
+
+    let values = get_elements(dict)
+        .values()
+        .map(|(_k, v)| v.clone())
+        .collect();
+    let values_list = vm.ctx.new_list(values);
+
+    let iter_obj = PyObject::new(
+        PyObjectPayload::Iterator {
+            position: Cell::new(0),
+            iterated_obj: values_list,
+        },
+        vm.ctx.iter_type(),
+    );
+
+    Ok(iter_obj)
+}
+
+fn dict_items(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(dict, Some(vm.ctx.dict_type()))]);
+
+    let items = get_elements(dict)
+        .values()
+        .map(|(k, v)| vm.ctx.new_tuple(vec![k.clone(), v.clone()]))
+        .collect();
+    let items_list = vm.ctx.new_list(items);
+
+    let iter_obj = PyObject::new(
+        PyObjectPayload::Iterator {
+            position: Cell::new(0),
+            iterated_obj: items_list,
+        },
+        vm.ctx.iter_type(),
+    );
+
+    Ok(iter_obj)
+}
+
 fn dict_setitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
@@ -336,4 +376,7 @@ pub fn init(context: &PyContext) {
         context.new_rustfunc(dict_setitem),
     );
     context.set_attr(&dict_type, "clear", context.new_rustfunc(dict_clear));
+    context.set_attr(&dict_type, "values", context.new_rustfunc(dict_values));
+    context.set_attr(&dict_type, "items", context.new_rustfunc(dict_items));
+    context.set_attr(&dict_type, "keys", context.new_rustfunc(dict_iter));
 }