defaultdict rust impl

ShaharNaveh · ShaharNaveh · commit 4857fbbf164c · 2026-06-19T16:39:29.000+03:00
diff --git a/Lib/collections/__init__.py b/Lib/collections/__init__.py
@@ -57,8 +57,7 @@
 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
diff --git 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<PyObjectRef>,
@@ -406,7 +406,7 @@ impl PyDict {
     }
 
     #[pymethod]
-    fn update(
+    pub(crate) fn update(
         &self,
         dict_obj: OptionalArg<PyObjectRef>,
         kwargs: KwArgs,
diff --git 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,200 @@ 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.read().clone();
+            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.read().clone();
+
+            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 factory_tuple = match &*zelf.default_factory.read() {
+                Some(f) => vm.ctx.new_tuple(vec![f.clone()]),
+                None => vm.ctx.new_tuple(vec![]),
+            };
+
+            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 {
+            Ok(if let Some(zelf) = lhs.downcast_ref::<Self>() {
+                if !rhs.fast_isinstance(vm.ctx.types.dict_type) {
+                    vm.ctx.not_implemented.clone().into()
+                } else {
+                    let default_factory = zelf.default_factory.read().clone();
+                    let dict = zelf.dict.copy();
+
+                    dict.update(rhs.into(), KwArgs::default(), vm)?;
+
+                    Self {
+                        dict,
+                        default_factory: PyRwLock::new(default_factory),
+                    }
+                    .to_pyobject(vm)
+                }
+            } else if let Some(zelf) = rhs.downcast_ref::<Self>() {
+                let default_factory = zelf.default_factory.read().clone();
+                if let Some(dict) = lhs.downcast_ref::<PyDict>() {
+                    let dict = dict.copy();
+                    dict.update(rhs.into(), KwArgs::default(), vm)?;
+
+                    Self {
+                        dict,
+                        default_factory: PyRwLock::new(default_factory),
+                    }
+                    .to_pyobject(vm)
+                } else {
+                    vm.ctx.not_implemented.clone().into()
+                }
+            } else {
+                return Err(vm.new_type_error(format!(
+                    "unsupported operand type(s) for |: '{}' and '{}'",
+                    lhs.class().name(),
+                    rhs.class().name()
+                )));
+            })
+        }
+    }
+
+    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
+        }
+    }
 }
