Skip to content

Commit 39e6f99

Browse files
morealyouknowone
authored andcommitted
Implement _collections.OrderedDict
Extracted from RustPython#6799: native OrderedDict as a dict subclass, iterator mutation tracking, and mapping-view set operations.
1 parent 2b8c095 commit 39e6f99

5 files changed

Lines changed: 1772 additions & 11 deletions

File tree

‎crates/vm/src/builtins/dict.rs‎

Lines changed: 70 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use super::{
44
};
55
use crate::common::lock::LazyLock;
66
use crate::object::{Traverse, TraverseFn};
7+
use crate::stdlib::PyOrderedDictItems;
78
use crate::{
89
AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
910
TryFromObject, atomic_func,
@@ -213,7 +214,7 @@ impl PyDict {
213214
Ok((key.clone(), value.clone()))
214215
}
215216

216-
fn update_sequence_pair(
217+
pub(crate) fn update_sequence_pair(
217218
element: PyObjectRef,
218219
index: usize,
219220
vm: &VirtualMachine,
@@ -496,8 +497,9 @@ impl PyDict {
496497
}
497498

498499
fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
499-
let other_dict: Result<PyDictRef, _> = other.downcast();
500-
if let Ok(other) = other_dict {
500+
// Accept dict subclasses that inherit this implementation. Subclasses with
501+
// their own reflected slot, such as OrderedDict, take precedence in dispatch.
502+
if let Ok(other) = other.downcast::<Self>() {
501503
let self_cp = self.copy();
502504
self_cp.merge_dict(other, true, vm)?;
503505
return Ok(self_cp.into_pyobject(vm));
@@ -542,7 +544,7 @@ impl PyDict {
542544

543545
#[pyclass]
544546
impl Py<PyDict> {
545-
fn inner_cmp(
547+
pub(crate) fn inner_cmp(
546548
&self,
547549
other: &Self,
548550
op: PyComparisonOp,
@@ -1675,7 +1677,7 @@ impl AsNumber for PyDictItems {
16751677
static AS_NUMBER: PyNumberMethods = PyNumberMethods {
16761678
subtract: Some(set_inner_number_subtract),
16771679
and: Some(set_inner_number_and),
1678-
xor: Some(set_inner_number_xor),
1680+
xor: Some(dict_item_view_number_xor),
16791681
or: Some(set_inner_number_or),
16801682
..PyNumberMethods::NOT_IMPLEMENTED
16811683
};
@@ -1716,19 +1718,77 @@ where
17161718
Ok(PySet { inner: f(a, b)? }.into_pyobject(vm))
17171719
}
17181720

1719-
fn set_inner_number_subtract(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1721+
pub(crate) fn set_inner_number_subtract(
1722+
a: &PyObject,
1723+
b: &PyObject,
1724+
vm: &VirtualMachine,
1725+
) -> PyResult {
17201726
set_inner_number_op(a, b, |a, b| a.difference(b, vm), vm)
17211727
}
17221728

1723-
fn set_inner_number_and(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1724-
set_inner_number_op(a, b, |a, b| a.intersection(b, vm), vm)
1729+
pub(crate) fn set_inner_number_and(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1730+
let a_is_view =
1731+
a.downcast_ref::<PyDictKeys>().is_some() || a.downcast_ref::<PyDictItems>().is_some();
1732+
let (view, other) = if a_is_view { (a, b) } else { (b, a) };
1733+
set_view_number_and(view, other, vm)
17251734
}
17261735

1727-
fn set_inner_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1736+
pub(crate) fn set_view_number_and(
1737+
view: &PyObject,
1738+
other: &PyObject,
1739+
vm: &VirtualMachine,
1740+
) -> PyResult {
1741+
let other = ArgIterable::<PyObjectRef>::try_from_object(vm, other.to_owned())?;
1742+
let result = PySet::default();
1743+
for item in other.iter(vm)? {
1744+
let item = item?;
1745+
if view.sequence_unchecked().contains(&item, vm)? {
1746+
result.add(item, vm)?;
1747+
}
1748+
}
1749+
Ok(result.into_pyobject(vm))
1750+
}
1751+
1752+
pub(crate) fn set_inner_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
17281753
set_inner_number_op(a, b, |a, b| a.symmetric_difference(b, vm), vm)
17291754
}
17301755

1731-
fn set_inner_number_or(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1756+
pub(crate) fn set_item_view_number_xor(
1757+
a: &PyObject,
1758+
b: &PyObject,
1759+
both_are_item_views: bool,
1760+
vm: &VirtualMachine,
1761+
) -> PyResult {
1762+
if !both_are_item_views {
1763+
return set_inner_number_xor(a, b, vm);
1764+
}
1765+
let result = PySet::default();
1766+
let a_iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, a.to_owned())?;
1767+
for item in a_iterable.iter(vm)? {
1768+
let item = item?;
1769+
if !b.sequence_unchecked().contains(&item, vm)? {
1770+
result.add(item, vm)?;
1771+
}
1772+
}
1773+
let b_iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, b.to_owned())?;
1774+
for item in b_iterable.iter(vm)? {
1775+
let item = item?;
1776+
if !a.sequence_unchecked().contains(&item, vm)? {
1777+
result.add(item, vm)?;
1778+
}
1779+
}
1780+
Ok(result.into_pyobject(vm))
1781+
}
1782+
1783+
fn dict_item_view_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1784+
let is_native_item_view = |obj: &PyObject| {
1785+
obj.downcast_ref::<PyDictItems>().is_some()
1786+
|| obj.downcast_ref::<PyOrderedDictItems>().is_some()
1787+
};
1788+
set_item_view_number_xor(a, b, is_native_item_view(a) && is_native_item_view(b), vm)
1789+
}
1790+
1791+
pub(crate) fn set_inner_number_or(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
17321792
set_inner_number_op(a, b, |a, b| a.union(b, vm), vm)
17331793
}
17341794

0 commit comments

Comments
 (0)