diff --git a/Lib/collections/__init__.py b/Lib/collections/__init__.py index 3d3bbd7a3..803de0c67 100644 --- a/Lib/collections/__init__.py +++ b/Lib/collections/__init__.py @@ -57,8 +57,7 @@ except ImportError: try: from _collections import defaultdict except ImportError: - # TODO: RUSTPYTHON - implement defaultdict in Rust - from ._defaultdict import defaultdict + pass heapq = None # Lazily imported diff --git a/Lib/collections/_defaultdict.py b/Lib/collections/_defaultdict.py deleted file mode 100644 index cb9d403c8..000000000 --- a/Lib/collections/_defaultdict.py +++ /dev/null @@ -1,62 +0,0 @@ -from reprlib import recursive_repr as _recursive_repr - -class defaultdict(dict): - def __init__(self, *args, **kwargs): - if len(args) >= 1: - default_factory = args[0] - if default_factory is not None and not callable(default_factory): - raise TypeError("first argument must be callable or None") - args = args[1:] - else: - default_factory = None - super().__init__(*args, **kwargs) - self.default_factory = default_factory - - def __missing__(self, key): - if self.default_factory is not None: - val = self.default_factory() - else: - raise KeyError(key) - # CPython parity: a recursive __missing__ via factory() may have - # already populated key; preserve that value instead of overwriting. - if key in self: - return self[key] - self[key] = val - return val - - @_recursive_repr() - def __repr_factory(factory): - return repr(factory) - - def __repr__(self): - return f"{type(self).__name__}({defaultdict.__repr_factory(self.default_factory)}, {dict.__repr__(self)})" - - def copy(self): - return type(self)(self.default_factory, self) - - __copy__ = copy - - def __reduce__(self): - if self.default_factory is not None: - args = self.default_factory, - else: - args = () - return type(self), args, None, None, iter(self.items()) - - def __or__(self, other): - if not isinstance(other, dict): - return NotImplemented - - new = defaultdict(self.default_factory, self) - new.update(other) - return new - - def __ror__(self, other): - if not isinstance(other, dict): - return NotImplemented - - new = defaultdict(self.default_factory, other) - new.update(self) - return new - -defaultdict.__module__ = 'collections' diff --git a/Lib/test/test_dataclasses/__init__.py b/Lib/test/test_dataclasses/__init__.py index 1a6dc850a..96f421832 100644 --- a/Lib/test/test_dataclasses/__init__.py +++ b/Lib/test/test_dataclasses/__init__.py @@ -1795,7 +1795,6 @@ class TestCase(unittest.TestCase): self.assertIsNot(d['f'], t) self.assertEqual(d['f'].my_a(), 6) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_helper_asdict_defaultdict(self): # Ensure asdict() does not throw exceptions when a # defaultdict is a member of a dataclass @@ -1938,7 +1937,6 @@ class TestCase(unittest.TestCase): t = astuple(c, tuple_factory=list) self.assertEqual(t, ['outer', T(1, ['inner', T(11, 12, 13)], 2)]) - @unittest.expectedFailure # TODO: RUSTPYTHON def test_helper_astuple_defaultdict(self): # Ensure astuple() does not throw exceptions when a # defaultdict is a member of a dataclass diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs index db97045f0..06791410e 100644 --- a/crates/vm/src/builtins/dict.rs +++ b/crates/vm/src/builtins/dict.rs @@ -387,7 +387,7 @@ impl PyDict { } #[pymethod] - fn setdefault( + pub(crate) fn setdefault( &self, key: PyObjectRef, default: OptionalArg, @@ -406,7 +406,7 @@ impl PyDict { } #[pymethod] - fn update( + pub(crate) fn update( &self, dict_obj: OptionalArg, kwargs: KwArgs, diff --git a/crates/vm/src/stdlib/_collections.rs b/crates/vm/src/stdlib/_collections.rs index 4f580ab5d..c7cce5c73 100644 --- a/crates/vm/src/stdlib/_collections.rs +++ b/crates/vm/src/stdlib/_collections.rs @@ -7,18 +7,19 @@ mod _collections { atomic_func, builtins::{ IterStatus::{Active, Exhausted}, - PositionIterInternal, PyGenericAlias, PyInt, PyStr, PyType, PyTypeRef, + PositionIterInternal, PyDict, PyGenericAlias, PyInt, PyStr, PyType, PyTypeRef, }, common::lock::{PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard}, - function::{KwArgs, OptionalArg, PyComparisonValue}, + convert::ToPyObject, + function::{FuncArgs, KwArgs, OptionalArg, PyComparisonValue}, iter::PyExactSizeIterator, - protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods}, + protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods}, recursion::ReprGuard, sequence::{MutObjectSequenceOp, OptionalRangeArgs}, sliceable::SequenceIndexOp, types::{ - AsNumber, AsSequence, Comparable, Constructor, DefaultConstructor, Initializer, - IterNext, Iterable, PyComparisonOp, Representable, SelfIter, + AsMapping, AsNumber, AsSequence, Comparable, Constructor, DefaultConstructor, + Initializer, IterNext, Iterable, PyComparisonOp, Representable, SelfIter, }, utils::collection_repr, }; @@ -746,4 +747,194 @@ mod _collections { }) } } + + #[pyattr] + #[pyclass( + module = "collections", + name = "defaultdict", + base = PyDict, + unhashable = true + )] + #[derive(Debug, Default)] + struct PyDefaultDict { + dict: PyDict, + default_factory: PyRwLock>, + } + + #[pyclass( + with(AsMapping, AsNumber, Constructor, Initializer, Representable), + flags(BASETYPE, MAPPING, HAS_DICT) + )] + impl PyDefaultDict { + #[pygetset] + fn default_factory(&self) -> Option { + self.default_factory.read().clone() + } + + #[pygetset(name = "default_factory", setter)] + fn default_factory_setter(&self, value: PyObjectRef, vm: &VirtualMachine) { + *self.default_factory.write() = if value.is(&vm.ctx.none()) { + None + } else { + Some(value) + }; + } + + #[pymethod] + fn __missing__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let factory = self.default_factory(); + + if let Some(f) = factory { + let value = f.call((), vm)?; + self.dict.setdefault(key, value.into(), vm) + } else { + Err(vm.new_key_error(key)) + } + } + + #[pymethod] + #[pymethod(name = "__copy__")] + fn copy(&self) -> Self { + let default_factory = self.default_factory(); + + Self { + dict: self.dict.copy(), + default_factory: PyRwLock::new(default_factory), + } + } + + #[pymethod] + fn __reduce__(zelf: PyRef, vm: &VirtualMachine) -> PyResult { + let cls = zelf.class().to_owned(); + + let default_factory = zelf.default_factory(); + let factory_tuple_elements = + default_factory.map_or_else(Vec::new, |factory| vec![factory]); + let factory_tuple = vm.ctx.new_tuple(factory_tuple_elements); + + let items_fn = zelf.as_object().get_attr("items", vm)?; + let items_iter = items_fn.call((), vm)?; + let iter = items_iter.get_iter(vm)?; + let none = vm.ctx.none(); + + Ok(vm + .ctx + .new_tuple(vec![ + cls.into(), + factory_tuple.into(), + none.clone(), + none, + iter.into(), + ]) + .into()) + } + } + + impl PyDefaultDict { + fn __or__(lhs: PyObjectRef, rhs: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let not_implemented = || Ok(vm.ctx.not_implemented.clone().into()); + + let (default_factory, dict) = if let Some(zelf) = lhs.downcast_ref::() { + if !rhs.fast_isinstance(vm.ctx.types.dict_type) { + return not_implemented(); + } + + (zelf.default_factory(), zelf.dict.copy()) + } else if let Some(zelf) = rhs.downcast_ref::() { + let Some(dict) = lhs.downcast_ref::() else { + return not_implemented(); + }; + + (zelf.default_factory(), dict.copy()) + } else { + return Err(vm.new_type_error(format!( + "unsupported operand type(s) for |: '{}' and '{}'", + lhs.class().name(), + rhs.class().name() + ))); + }; + + dict.update(rhs.into(), KwArgs::default(), vm)?; + + Ok(Self { + dict, + default_factory: PyRwLock::new(default_factory), + } + .to_pyobject(vm)) + } + } + + impl DefaultConstructor for PyDefaultDict {} + + impl Initializer for PyDefaultDict { + type Args = FuncArgs; + + fn init(zelf: PyRef, mut args: Self::Args, vm: &VirtualMachine) -> PyResult<()> { + let default_factory = args.take_positional().map_or(Ok(None), |factory| { + let is_none = factory.is(&vm.ctx.none()); + + if !is_none && !factory.is_callable() { + Err(vm.new_type_error("first argument must be callable or None")) + } else if is_none { + Ok(None) + } else { + Ok(Some(factory)) + } + })?; + + *zelf.default_factory.write() = default_factory; + + zelf.dict.update( + OptionalArg::from_option(args.take_positional()), + args.kwargs, + vm, + )?; + + Ok(()) + } + } + + impl Representable for PyDefaultDict { + fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let default_factory = zelf.default_factory.read(); + + let factory_repr = match default_factory.as_ref() { + Some(factory) => { + if let Some(_guard) = ReprGuard::enter(vm, factory) { + factory.repr(vm)?.to_string() + } else { + String::from("...") + } + } + None => String::from("None"), + }; + + let dict_repr = Representable::repr(&zelf.dict.copy().into_ref(&vm.ctx), vm)?; + + Ok(format!( + "{}({}, {})", + zelf.class().name(), + factory_repr, + dict_repr + )) + } + } + + impl AsMapping for PyDefaultDict { + fn as_mapping() -> &'static PyMappingMethods { + PyDict::as_mapping() + } + } + + impl AsNumber for PyDefaultDict { + fn as_number() -> &'static PyNumberMethods { + static AS_NUMBER: PyNumberMethods = PyNumberMethods { + or: Some(|a, b, vm| { + PyDefaultDict::__or__(a.to_pyobject(vm), b.to_pyobject(vm), vm) + }), + ..PyNumberMethods::NOT_IMPLEMENTED + }; + &AS_NUMBER + } + } }