diff --git a/crates/vm/src/stdlib/_collections/ordered_dict.rs b/crates/vm/src/stdlib/_collections/ordered_dict.rs index 4ba9777fe..8c3079c3b 100644 --- a/crates/vm/src/stdlib/_collections/ordered_dict.rs +++ b/crates/vm/src/stdlib/_collections/ordered_dict.rs @@ -7,7 +7,7 @@ pub(crate) mod ordered_dict { AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine, atomic_func, builtins::{ - PyDict, PyGenericAlias, PyMappingProxy, PyTuple, PyTypeRef, + PyDict, PyGenericAlias, PyMappingProxy, PyStrRef, PyTuple, PyTypeRef, dict::{ PyDictItems, set_inner_number_or, set_inner_number_subtract, set_inner_number_xor, set_item_view_number_xor, set_view_number_and, @@ -27,6 +27,7 @@ pub(crate) mod ordered_dict { Initializer, IterNext, Iterable, PyComparisonOp, Representable, SelfIter, }, }; + use rustpython_common::wtf8::Wtf8Buf; use std::collections::HashMap; const ODICT_ITER_REVERSED: u8 = 1; @@ -977,6 +978,29 @@ pub(crate) mod ordered_dict { set_item_view_number_xor(a, b, is_item_view(a) && is_item_view(b), vm) } + fn ordered_view_repr( + view: &PyObject, + od: &PyRef, + kind: u8, + name: &str, + vm: &VirtualMachine, + ) -> PyResult { + let Some(_guard) = ReprGuard::enter(vm, view) else { + return Ok(vm.ctx.intern_str("...").to_owned()); + }; + // Snapshot the view before calling user-defined repr methods, which + // can mutate the underlying ordered dictionary. + let iterator = PyODictIter::new(od.clone(), kind).to_pyobject(vm); + let items = PyIter::try_from_object(vm, iterator)? + .iter::(vm)? + .collect::>>()?; + let list_repr = vm.ctx.new_list(items).as_object().repr(vm)?; + let mut result = Wtf8Buf::from(format!("{name}(")); + result.push_wtf8(list_repr.as_wtf8()); + result.push_str(")"); + Ok(vm.ctx.new_str(result)) + } + #[pyattr] #[pyclass(name = "odict_keys", module = "builtins", unhashable = true, traverse)] #[derive(Debug, PyPayload)] @@ -984,7 +1008,7 @@ pub(crate) mod ordered_dict { od: PyRef, } - #[pyclass(with(Iterable, Comparable, AsMapping, AsSequence, AsNumber))] + #[pyclass(with(Iterable, Comparable, AsMapping, AsSequence, AsNumber, Representable))] impl Py { #[pymethod] fn __reversed__(&self, vm: &VirtualMachine) -> PyObjectRef { @@ -1007,6 +1031,22 @@ pub(crate) mod ordered_dict { } } + impl Representable for PyODictKeys { + fn repr(zelf: &Py, vm: &VirtualMachine) -> PyResult { + ordered_view_repr( + zelf.as_object(), + &zelf.od, + ODICT_ITER_KEYS, + "odict_keys", + vm, + ) + } + + fn repr_str(_zelf: &Py, _vm: &VirtualMachine) -> PyResult { + unreachable!("use repr instead") + } + } + impl AsMapping for PyODictKeys { fn as_mapping() -> &'static PyMappingMethods { static AS_MAPPING: PyMappingMethods = PyMappingMethods { @@ -1077,7 +1117,7 @@ pub(crate) mod ordered_dict { od: PyRef, } - #[pyclass(with(Iterable, Comparable, AsMapping))] + #[pyclass(with(Iterable, Comparable, AsMapping, Representable))] impl Py { #[pymethod] fn __reversed__(&self, vm: &VirtualMachine) -> PyObjectRef { @@ -1091,6 +1131,22 @@ pub(crate) mod ordered_dict { } } + impl Representable for PyODictValues { + fn repr(zelf: &Py, vm: &VirtualMachine) -> PyResult { + ordered_view_repr( + zelf.as_object(), + &zelf.od, + ODICT_ITER_VALUES, + "odict_values", + vm, + ) + } + + fn repr_str(_zelf: &Py, _vm: &VirtualMachine) -> PyResult { + unreachable!("use repr instead") + } + } + impl AsMapping for PyODictValues { fn as_mapping() -> &'static PyMappingMethods { static AS_MAPPING: PyMappingMethods = PyMappingMethods { @@ -1128,7 +1184,7 @@ pub(crate) mod ordered_dict { od: PyRef, } - #[pyclass(with(Iterable, Comparable, AsMapping, AsSequence, AsNumber))] + #[pyclass(with(Iterable, Comparable, AsMapping, AsSequence, AsNumber, Representable))] impl Py { #[pymethod] fn __reversed__(&self, vm: &VirtualMachine) -> PyObjectRef { @@ -1163,6 +1219,22 @@ pub(crate) mod ordered_dict { } } + impl Representable for PyOrderedDictItems { + fn repr(zelf: &Py, vm: &VirtualMachine) -> PyResult { + ordered_view_repr( + zelf.as_object(), + &zelf.od, + ODICT_ITER_ITEMS, + "odict_items", + vm, + ) + } + + fn repr_str(_zelf: &Py, _vm: &VirtualMachine) -> PyResult { + unreachable!("use repr instead") + } + } + impl AsMapping for PyOrderedDictItems { fn as_mapping() -> &'static PyMappingMethods { static AS_MAPPING: PyMappingMethods = PyMappingMethods { diff --git a/extra_tests/snippets/stdlib_collections.py b/extra_tests/snippets/stdlib_collections.py index 2da8df87c..be9e7a99c 100644 --- a/extra_tests/snippets/stdlib_collections.py +++ b/extra_tests/snippets/stdlib_collections.py @@ -91,6 +91,60 @@ assert Counter("abracadabra").most_common(1) == [("a", 5)] ordered = OrderedDict() _count_elements(ordered, "bab") assert list(ordered.items()) == [("b", 2), ("a", 1)] + + +# OrderedDict views show their contents in insertion order, including recursion. +def check_ordered_dict_view_repr(): + ordered = OrderedDict([("b", 2), ("a", 1)]) + for method, contents in ( + ("keys", "['b', 'a']"), + ("values", "[2, 1]"), + ("items", "[('b', 2), ('a', 1)]"), + ): + assert repr(getattr(OrderedDict(), method)()) == "odict_" + method + "([])" + assert ( + repr(getattr(ordered, method)()) == "odict_" + method + "(" + contents + ")" + ) + + recursive = OrderedDict() + values = recursive.values() + items = recursive.items() + recursive["self"] = values + assert repr(values) == "odict_values([...])" + recursive["self"] = items + assert repr(items) == "odict_items([('self', ...)])" + + # Repr must snapshot the contents before arbitrary repr code mutates them. + for method in ("keys", "values", "items"): + + class MutatingRepr: + def __repr__(self): + ordered.clear() + return "mutated" + + item = MutatingRepr() + ordered = OrderedDict([(item, item), ("tail", "last")]) + expected = { + "keys": "odict_keys([mutated, 'tail'])", + "values": "odict_values([mutated, 'last'])", + "items": "odict_items([(mutated, mutated), ('tail', 'last')])", + } + assert repr(getattr(ordered, method)()) == expected[method] + assert not ordered + + class BadRepr: + def __repr__(self): + raise ValueError("view repr") + + ordered = OrderedDict([(BadRepr(), BadRepr())]) + for method in ("keys", "values", "items"): + with assert_raises(ValueError) as caught: + repr(getattr(ordered, method)()) + assert str(caught.exception) == "view repr" + + +check_ordered_dict_view_repr() + factory = defaultdict(lambda: 100) _count_elements(factory, "a") assert factory == {"a": 1}