1use crate::{types::PyComparisonOp, vm::VirtualMachine, PyObjectRef, PyResult};
2use itertools::Itertools;
3
4pub trait PyExactSizeIterator<'a>: ExactSizeIterator<Item = &'a PyObjectRef> + Sized {
5 fn eq(self, other: impl PyExactSizeIterator<'a>, vm: &VirtualMachine) -> PyResult<bool> {
6 let lhs = self;
7 let rhs = other;
8 if lhs.len() != rhs.len() {
9 return Ok(false);
10 }
11 for (a, b) in lhs.zip_eq(rhs) {
12 if !vm.identical_or_equal(a, b)? {
13 return Ok(false);
14 }
15 }
16 Ok(true)
17 }
18
19 fn richcompare(
20 self,
21 other: impl PyExactSizeIterator<'a>,
22 op: PyComparisonOp,
23 vm: &VirtualMachine,
24 ) -> PyResult<bool> {
25 let less = match op {
26 PyComparisonOp::Eq => return PyExactSizeIterator::eq(self, other, vm),
27 PyComparisonOp::Ne => return PyExactSizeIterator::eq(self, other, vm).map(|eq| !eq),
28 PyComparisonOp::Lt | PyComparisonOp::Le => true,
29 PyComparisonOp::Gt | PyComparisonOp::Ge => false,
30 };
31
32 let lhs = self;
33 let rhs = other;
34 let lhs_len = lhs.len();
35 let rhs_len = rhs.len();
36 for (a, b) in lhs.zip(rhs) {
37 if vm.bool_eq(a, b)? {
38 continue;
39 }
40 let ret = if less {
41 vm.bool_seq_lt(a, b)?
42 } else {
43 vm.bool_seq_gt(a, b)?
44 };
45 if let Some(v) = ret {
46 return Ok(v);
47 }
48 }
49 Ok(op.eval_ord(lhs_len.cmp(&rhs_len)))
50 }
51}
52
53impl<'a, T> PyExactSizeIterator<'a> for T where T: ExactSizeIterator<Item = &'a PyObjectRef> + Sized {}