Skip to content

Commit 7ad7465

Browse files
committed
Use PyIterReturn result-like tools
1 parent 2dea5f2 commit 7ad7465

File tree

9 files changed

+51
-67
lines changed

9 files changed

+51
-67
lines changed

vm/src/builtins/bytearray.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -844,10 +844,9 @@ impl IterNext for PyByteArrayIterator {
844844
fn next(zelf: &crate::PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
845845
zelf.internal.lock().next(|bytearray, pos| {
846846
let buf = bytearray.borrow_buf();
847-
Ok(match buf.get(pos) {
848-
Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x).into()),
849-
None => PyIterReturn::StopIteration(None),
850-
})
847+
Ok(PyIterReturn::from_result(
848+
buf.get(pos).map(|&x| vm.new_pyobj(x)).ok_or(None),
849+
))
851850
})
852851
}
853852
}

vm/src/builtins/bytes.rs

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -678,10 +678,13 @@ impl IterNextIterable for PyBytesIterator {}
678678
impl IterNext for PyBytesIterator {
679679
fn next(zelf: &crate::PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
680680
zelf.internal.lock().next(|bytes, pos| {
681-
Ok(match bytes.as_bytes().get(pos) {
682-
Some(&x) => PyIterReturn::Return(vm.ctx.new_int(x).into()),
683-
None => PyIterReturn::StopIteration(None),
684-
})
681+
Ok(PyIterReturn::from_result(
682+
bytes
683+
.as_bytes()
684+
.get(pos)
685+
.map(|&x| vm.new_pyobj(x))
686+
.ok_or(None),
687+
))
685688
})
686689
}
687690
}

vm/src/builtins/dict.rs

Lines changed: 14 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -105,14 +105,8 @@ impl PyDict {
105105
PyIterReturn::StopIteration(_) => break,
106106
};
107107
let elem_iter = element.get_iter(vm)?;
108-
let key = match elem_iter.next(vm)? {
109-
PyIterReturn::Return(obj) => obj,
110-
PyIterReturn::StopIteration(_) => return Err(err(vm)),
111-
};
112-
let value = match elem_iter.next(vm)? {
113-
PyIterReturn::Return(obj) => obj,
114-
PyIterReturn::StopIteration(_) => return Err(err(vm)),
115-
};
108+
let key = elem_iter.next(vm)?.into_result().map_err(|_| err(vm))?;
109+
let value = elem_iter.next(vm)?.into_result().map_err(|_| err(vm))?;
116110
if matches!(elem_iter.next(vm)?, PyIterReturn::Return(_)) {
117111
return Err(err(vm));
118112
}
@@ -814,7 +808,7 @@ macro_rules! dict_view {
814808
#[allow(clippy::redundant_closure_call)]
815809
fn next(zelf: &PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
816810
let mut internal = zelf.internal.lock();
817-
if let IterStatus::Active(dict) = &internal.status {
811+
let next = if let IterStatus::Active(dict) = &internal.status {
818812
if dict.entries.has_changed_size(&zelf.size) {
819813
internal.status = IterStatus::Exhausted;
820814
return Err(vm.new_runtime_error(
@@ -824,16 +818,17 @@ macro_rules! dict_view {
824818
match dict.entries.next_entry(internal.position) {
825819
Some((position, key, value)) => {
826820
internal.position = position;
827-
Ok(PyIterReturn::Return(($result_fn)(vm, key, value)))
821+
PyIterReturn::Return(($result_fn)(vm, key, value))
828822
}
829823
None => {
830824
internal.status = IterStatus::Exhausted;
831-
Ok(PyIterReturn::StopIteration(None))
825+
PyIterReturn::StopIteration(None)
832826
}
833827
}
834828
} else {
835-
Ok(PyIterReturn::StopIteration(None))
836-
}
829+
PyIterReturn::StopIteration(None)
830+
};
831+
Ok(next)
837832
}
838833
}
839834

@@ -875,7 +870,7 @@ macro_rules! dict_view {
875870
#[allow(clippy::redundant_closure_call)]
876871
fn next(zelf: &PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
877872
let mut internal = zelf.internal.lock();
878-
if let IterStatus::Active(dict) = &internal.status {
873+
let next = if let IterStatus::Active(dict) = &internal.status {
879874
if dict.entries.has_changed_size(&zelf.size) {
880875
internal.status = IterStatus::Exhausted;
881876
return Err(vm.new_runtime_error(
@@ -889,16 +884,17 @@ macro_rules! dict_view {
889884
} else {
890885
internal.position = position;
891886
}
892-
Ok(PyIterReturn::Return(($result_fn)(vm, key, value)))
887+
PyIterReturn::Return(($result_fn)(vm, key, value))
893888
}
894889
None => {
895890
internal.status = IterStatus::Exhausted;
896-
Ok(PyIterReturn::StopIteration(None))
891+
PyIterReturn::StopIteration(None)
897892
}
898893
}
899894
} else {
900-
Ok(PyIterReturn::StopIteration(None))
901-
}
895+
PyIterReturn::StopIteration(None)
896+
};
897+
Ok(next)
902898
}
903899
}
904900
};

vm/src/builtins/iter.rs

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -236,17 +236,18 @@ impl IterNextIterable for PyCallableIterator {}
236236
impl IterNext for PyCallableIterator {
237237
fn next(zelf: &crate::PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
238238
let status = zelf.status.upgradable_read();
239-
if let IterStatus::Active(callable) = &*status {
239+
let next = if let IterStatus::Active(callable) = &*status {
240240
let ret = callable.invoke((), vm)?;
241241
if vm.bool_eq(&ret, &zelf.sentinel)? {
242242
*PyRwLockUpgradableReadGuard::upgrade(status) = IterStatus::Exhausted;
243-
Ok(PyIterReturn::StopIteration(None))
243+
PyIterReturn::StopIteration(None)
244244
} else {
245-
Ok(PyIterReturn::Return(ret))
245+
PyIterReturn::Return(ret)
246246
}
247247
} else {
248-
Ok(PyIterReturn::StopIteration(None))
249-
}
248+
PyIterReturn::StopIteration(None)
249+
};
250+
Ok(next)
250251
}
251252
}
252253

vm/src/builtins/list.rs

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -668,10 +668,7 @@ impl IterNext for PyListIterator {
668668
fn next(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
669669
zelf.internal.lock().next(|list, pos| {
670670
let vec = list.borrow_vec();
671-
Ok(match vec.get(pos) {
672-
Some(x) => PyIterReturn::Return(x.clone()),
673-
None => PyIterReturn::StopIteration(None),
674-
})
671+
Ok(PyIterReturn::from_result(vec.get(pos).cloned().ok_or(None)))
675672
})
676673
}
677674
}
@@ -716,10 +713,7 @@ impl IterNext for PyListReverseIterator {
716713
fn next(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
717714
zelf.internal.lock().rev_next(|list, pos| {
718715
let vec = list.borrow_vec();
719-
Ok(match vec.get(pos) {
720-
Some(x) => PyIterReturn::Return(x.clone()),
721-
None => PyIterReturn::StopIteration(None),
722-
})
716+
Ok(PyIterReturn::from_result(vec.get(pos).cloned().ok_or(None)))
723717
})
724718
}
725719
}

vm/src/builtins/set.rs

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -920,24 +920,25 @@ impl IterNextIterable for PySetIterator {}
920920
impl IterNext for PySetIterator {
921921
fn next(zelf: &crate::PyObjectView<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
922922
let mut internal = zelf.internal.lock();
923-
if let IterStatus::Active(dict) = &internal.status {
923+
let next = if let IterStatus::Active(dict) = &internal.status {
924924
if dict.has_changed_size(&zelf.size) {
925925
internal.status = IterStatus::Exhausted;
926926
return Err(vm.new_runtime_error("set changed size during iteration".to_owned()));
927927
}
928928
match dict.next_entry(internal.position) {
929929
Some((position, key, _)) => {
930930
internal.position = position;
931-
Ok(PyIterReturn::Return(key))
931+
PyIterReturn::Return(key)
932932
}
933933
None => {
934934
internal.status = IterStatus::Exhausted;
935-
Ok(PyIterReturn::StopIteration(None))
935+
PyIterReturn::StopIteration(None)
936936
}
937937
}
938938
} else {
939-
Ok(PyIterReturn::StopIteration(None))
940-
}
939+
PyIterReturn::StopIteration(None)
940+
};
941+
Ok(next)
941942
}
942943
}
943944

vm/src/builtins/tuple.rs

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -412,11 +412,9 @@ impl IterNextIterable for PyTupleIterator {}
412412
impl IterNext for PyTupleIterator {
413413
fn next(zelf: &crate::PyObjectView<Self>, _vm: &VirtualMachine) -> PyResult<PyIterReturn> {
414414
zelf.internal.lock().next(|tuple, pos| {
415-
Ok(if let Some(ret) = tuple.as_slice().get(pos) {
416-
PyIterReturn::Return(ret.clone())
417-
} else {
418-
PyIterReturn::StopIteration(None)
419-
})
415+
Ok(PyIterReturn::from_result(
416+
tuple.as_slice().get(pos).cloned().ok_or(None),
417+
))
420418
})
421419
}
422420
}

vm/src/protocol/iter.rs

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -223,10 +223,7 @@ where
223223
fn next(&mut self) -> Option<Self::Item> {
224224
PyIter::new(self.obj.borrow())
225225
.next(self.vm)
226-
.map(|iret| match iret {
227-
PyIterReturn::Return(obj) => Some(obj),
228-
PyIterReturn::StopIteration(_) => None,
229-
})
226+
.map(|iret| iret.into_result().ok())
230227
.transpose()
231228
.map(|x| x.and_then(|obj| T::try_from_object(self.vm, obj)))
232229
}

vm/src/stdlib/collections.rs

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -634,10 +634,9 @@ mod _collections {
634634
return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned()));
635635
}
636636
let deque = deque.borrow_deque();
637-
Ok(match deque.get(pos) {
638-
Some(x) => PyIterReturn::Return(x.clone()),
639-
None => PyIterReturn::StopIteration(None),
640-
})
637+
Ok(PyIterReturn::from_result(
638+
deque.get(pos).cloned().ok_or(None),
639+
))
641640
})
642641
}
643642
}
@@ -701,16 +700,12 @@ mod _collections {
701700
return Err(vm.new_runtime_error("Deque mutated during iteration".to_owned()));
702701
}
703702
let deque = deque.borrow_deque();
704-
Ok(
705-
match deque
706-
.len()
707-
.checked_sub(pos + 1)
708-
.and_then(|pos| deque.get(pos))
709-
{
710-
Some(x) => PyIterReturn::Return(x.clone()),
711-
None => PyIterReturn::StopIteration(None),
712-
},
713-
)
703+
let r = deque
704+
.len()
705+
.checked_sub(pos + 1)
706+
.and_then(|pos| deque.get(pos))
707+
.cloned();
708+
Ok(PyIterReturn::from_result(r.ok_or(None)))
714709
})
715710
}
716711
}

0 commit comments

Comments
 (0)