Use raise_if_stop! macro where possible (#5938)

This commit is contained in:
Shahar Naveh authored and GitHub committed 2025-07-11 08:36:08 +09:00
1 parent 38837e587b
commit 01f15065fa
3 files changed
+15 -18

No files matched your search

+3 -4
View File
@@ -8,6 +8,7 @@ use crate::{
convert::ToPyObject,
function::OptionalArg,
protocol::{PyIter, PyIterReturn},
raise_if_stop,
types::{Constructor, IterNext, Iterable, SelfIter},
};
use malachite_bigint::BigInt;
@@ -73,12 +74,10 @@ impl Py<PyEnumerate> {
}
impl SelfIter for PyEnumerate {}
impl IterNext for PyEnumerate {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
let next_obj = match zelf.iterator.next(vm)? {
PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
PyIterReturn::Return(obj) => obj,
};
let next_obj = raise_if_stop!(zelf.iterator.next(vm)?);
let mut counter = zelf.counter.write();
let position = counter.clone();
*counter += 1;
+9 -10
View File
@@ -3,6 +3,7 @@ use crate::{
Context, Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
class::PyClassImpl,
protocol::{PyIter, PyIterReturn},
raise_if_stop,
types::{Constructor, IterNext, Iterable, SelfIter},
};
@@ -45,24 +46,22 @@ impl PyFilter {
}
impl SelfIter for PyFilter {}
impl IterNext for PyFilter {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
let predicate = &zelf.predicate;
loop {
let next_obj = match zelf.iterator.next(vm)? {
PyIterReturn::Return(obj) => obj,
PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
};
let next_obj = raise_if_stop!(zelf.iterator.next(vm)?);
let predicate_value = if vm.is_none(predicate) {
next_obj.clone()
} else {
// the predicate itself can raise StopIteration which does stop the filter
// iteration
match PyIterReturn::from_pyresult(predicate.call((next_obj.clone(),), vm), vm)? {
PyIterReturn::Return(obj) => obj,
PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
}
// the predicate itself can raise StopIteration which does stop the filter iteration
raise_if_stop!(PyIterReturn::from_pyresult(
predicate.call((next_obj.clone(),), vm),
vm
)?)
};
if predicate_value.try_to_bool(vm)? {
return Ok(PyIterReturn::Return(next_obj));
}
+3 -4
View File
@@ -5,6 +5,7 @@ use crate::{
class::PyClassImpl,
function::PosArgs,
protocol::{PyIter, PyIterReturn},
raise_if_stop,
types::{Constructor, IterNext, Iterable, SelfIter},
};
@@ -53,14 +54,12 @@ impl PyMap {
}
impl SelfIter for PyMap {}
impl IterNext for PyMap {
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
let mut next_objs = Vec::new();
for iterator in &zelf.iterators {
let item = match iterator.next(vm)? {
PyIterReturn::Return(obj) => obj,
PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)),
};
let item = raise_if_stop!(iterator.next(vm)?);
next_objs.push(item);
}