diff --git a/tests/snippets/dict.py b/tests/snippets/dict.py index 1e4acc3e3d..170a6492dc 100644 --- a/tests/snippets/dict.py +++ b/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']) diff --git a/vm/src/obj/objdict.rs b/vm/src/obj/objdict.rs index 03a307b4c3..189684eb05 100644 --- a/vm/src/obj/objdict.rs +++ b/vm/src/obj/objdict.rs @@ -252,8 +252,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); @@ -268,6 +268,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, @@ -345,4 +385,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)); }