Implement collections.defaultdict in rust (#8132)

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
 

Lib/collections/_defaultdict.py

@@ -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'

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

crates/vm/src/builtins/dict.rs

@@ -387,7 +387,7 @@ impl PyDict {
     }
 
     #[pymethod]
-    fn setdefault(
+    pub(crate) fn setdefault(
         &self,
         key: PyObjectRef,
         default: OptionalArg<PyObjectRef>,
@@ -406,7 +406,7 @@ impl PyDict {
     }
 
     #[pymethod]
-    fn update(
+    pub(crate) fn update(
         &self,
         dict_obj: OptionalArg<PyObjectRef>,
         kwargs: KwArgs,

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<Option<PyObjectRef>>,
+    }
+
+    #[pyclass(
+        with(AsMapping, AsNumber, Constructor, Initializer, Representable),
+        flags(BASETYPE, MAPPING, HAS_DICT)
+    )]
+    impl PyDefaultDict {
+        #[pygetset]
+        fn default_factory(&self) -> Option<PyObjectRef> {
+            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<Self>, 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::<Self>() {
+                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::<Self>() {
+                let Some(dict) = lhs.downcast_ref::<PyDict>() 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<Self>, 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<Self>, vm: &VirtualMachine) -> PyResult<String> {
+            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
+        }
+    }
 }