Skip to content

Commit e9e8847

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 05f1386 commit e9e8847

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,
@@ -1667,7 +1669,7 @@ impl AsNumber for PyDictItems {
16671669
static AS_NUMBER: PyNumberMethods = PyNumberMethods {
16681670
subtract: Some(set_inner_number_subtract),
16691671
and: Some(set_inner_number_and),
1670-
xor: Some(set_inner_number_xor),
1672+
xor: Some(dict_item_view_number_xor),
16711673
or: Some(set_inner_number_or),
16721674
..PyNumberMethods::NOT_IMPLEMENTED
16731675
};
@@ -1708,19 +1710,77 @@ where
17081710
Ok(PySet { inner: f(a, b)? }.into_pyobject(vm))
17091711
}
17101712

1711-
fn set_inner_number_subtract(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1713+
pub(crate) fn set_inner_number_subtract(
1714+
a: &PyObject,
1715+
b: &PyObject,
1716+
vm: &VirtualMachine,
1717+
) -> PyResult {
17121718
set_inner_number_op(a, b, |a, b| a.difference(b, vm), vm)
17131719
}
17141720

1715-
fn set_inner_number_and(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1716-
set_inner_number_op(a, b, |a, b| a.intersection(b, vm), vm)
1721+
pub(crate) fn set_inner_number_and(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1722+
let a_is_view =
1723+
a.downcast_ref::<PyDictKeys>().is_some() || a.downcast_ref::<PyDictItems>().is_some();
1724+
let (view, other) = if a_is_view { (a, b) } else { (b, a) };
1725+
set_view_number_and(view, other, vm)
17171726
}
17181727

1719-
fn set_inner_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1728+
pub(crate) fn set_view_number_and(
1729+
view: &PyObject,
1730+
other: &PyObject,
1731+
vm: &VirtualMachine,
1732+
) -> PyResult {
1733+
let other = ArgIterable::<PyObjectRef>::try_from_object(vm, other.to_owned())?;
1734+
let result = PySet::default();
1735+
for item in other.iter(vm)? {
1736+
let item = item?;
1737+
if view.sequence_unchecked().contains(&item, vm)? {
1738+
result.add(item, vm)?;
1739+
}
1740+
}
1741+
Ok(result.into_pyobject(vm))
1742+
}
1743+
1744+
pub(crate) fn set_inner_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
17201745
set_inner_number_op(a, b, |a, b| a.symmetric_difference(b, vm), vm)
17211746
}
17221747

1723-
fn set_inner_number_or(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1748+
pub(crate) fn set_item_view_number_xor(
1749+
a: &PyObject,
1750+
b: &PyObject,
1751+
both_are_item_views: bool,
1752+
vm: &VirtualMachine,
1753+
) -> PyResult {
1754+
if !both_are_item_views {
1755+
return set_inner_number_xor(a, b, vm);
1756+
}
1757+
let result = PySet::default();
1758+
let a_iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, a.to_owned())?;
1759+
for item in a_iterable.iter(vm)? {
1760+
let item = item?;
1761+
if !b.sequence_unchecked().contains(&item, vm)? {
1762+
result.add(item, vm)?;
1763+
}
1764+
}
1765+
let b_iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, b.to_owned())?;
1766+
for item in b_iterable.iter(vm)? {
1767+
let item = item?;
1768+
if !a.sequence_unchecked().contains(&item, vm)? {
1769+
result.add(item, vm)?;
1770+
}
1771+
}
1772+
Ok(result.into_pyobject(vm))
1773+
}
1774+
1775+
fn dict_item_view_number_xor(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
1776+
let is_native_item_view = |obj: &PyObject| {
1777+
obj.downcast_ref::<PyDictItems>().is_some()
1778+
|| obj.downcast_ref::<PyOrderedDictItems>().is_some()
1779+
};
1780+
set_item_view_number_xor(a, b, is_native_item_view(a) && is_native_item_view(b), vm)
1781+
}
1782+
1783+
pub(crate) fn set_inner_number_or(a: &PyObject, b: &PyObject, vm: &VirtualMachine) -> PyResult {
17241784
set_inner_number_op(a, b, |a, b| a.union(b, vm), vm)
17251785
}
17261786

0 commit comments

Comments
 (0)