From e7e7ddbfcc8aeafdec7942e50382fe0c84990002 Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 12:26:32 +0300 Subject: [PATCH 1/9] Replace struct name with Self --- vm/src/stdlib/itertools.rs | 57 ++++++++++++++++++++------------------ 1 file changed, 30 insertions(+), 27 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index efe3368129..e511548ade 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -41,7 +41,7 @@ mod decl { #[pyslot] fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult { let args_list = PyList::from(args.args); - PyItertoolsChain { + Self { source: PyRwLock::new(Some(args_list.to_pyobject(vm).get_iter(vm)?)), active: PyRwLock::new(None), } @@ -91,17 +91,18 @@ mod decl { fn __setstate__(zelf: PyRef, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> { let args = state.as_slice(); if args.is_empty() { - let msg = String::from("function takes at least 1 arguments (0 given)"); - return Err(vm.new_type_error(msg)); + return Err(vm.new_type_error("function takes at least 1 arguments (0 given)")); } if args.len() > 2 { - let msg = format!("function takes at most 2 arguments ({} given)", args.len()); - return Err(vm.new_type_error(msg)); + return Err(vm.new_type_error(format!( + "function takes at most 2 arguments ({} given)", + args.len() + ))); } let source = &args[0]; if args.len() == 1 { if !PyIter::check(source.as_ref()) { - return Err(vm.new_type_error(String::from("Arguments must be iterators."))); + return Err(vm.new_type_error("Arguments must be iterators.")); } *zelf.source.write() = source.to_owned().try_into_value(vm)?; return Ok(()); @@ -109,7 +110,7 @@ mod decl { let active = &args[1]; if !PyIter::check(source.as_ref()) || !PyIter::check(active.as_ref()) { - return Err(vm.new_type_error(String::from("Arguments must be iterators."))); + return Err(vm.new_type_error("Arguments must be iterators.")); } let mut source_lock = zelf.source.write(); let mut active_lock = zelf.active.write(); @@ -387,7 +388,7 @@ mod decl { } None => None, }; - PyItertoolsRepeat { object, times } + Self { object, times } .into_ref_with_type(vm, cls) .map(Into::into) } @@ -466,7 +467,7 @@ mod decl { Self::Args { function, iterable }: Self::Args, vm: &VirtualMachine, ) -> PyResult { - PyItertoolsStarmap { function, iterable } + Self { function, iterable } .into_ref_with_type(vm, cls) .map(Into::into) } @@ -527,7 +528,7 @@ mod decl { }: Self::Args, vm: &VirtualMachine, ) -> PyResult { - PyItertoolsTakewhile { + Self { predicate, iterable, stop_flag: AtomicCell::new(false), @@ -614,7 +615,7 @@ mod decl { }: Self::Args, vm: &VirtualMachine, ) -> PyResult { - PyItertoolsDropwhile { + Self { predicate, iterable, start_flag: AtomicCell::new(false), @@ -634,6 +635,7 @@ mod decl { (zelf.start_flag.load() as _), ) } + #[pymethod] fn __setstate__( zelf: PyRef, @@ -734,7 +736,7 @@ mod decl { Self::Args { iterable, key }: Self::Args, vm: &VirtualMachine, ) -> PyResult { - PyItertoolsGroupBy { + Self { iterable, key_func: key.flatten(), state: PyMutex::new(GroupByState { @@ -952,7 +954,7 @@ mod decl { let iter = iter.get_iter(vm)?; - PyItertoolsIslice { + Self { iterable: iter, cur: AtomicCell::new(0), next: AtomicCell::new(start), @@ -980,14 +982,16 @@ mod decl { fn __setstate__(zelf: PyRef, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> { let args = state.as_slice(); if args.len() != 1 { - let msg = format!("function takes exactly 1 argument ({} given)", args.len()); - return Err(vm.new_type_error(msg)); + return Err(vm.new_type_error(format!( + "function takes exactly 1 argument ({} given)", + args.len() + ))); } let cur = &args[0]; if let Ok(cur) = cur.try_to_value(vm) { zelf.cur.store(cur); } else { - return Err(vm.new_type_error(String::from("Argument must be usize."))); + return Err(vm.new_type_error("Argument must be usize.")); } Ok(()) } @@ -1049,7 +1053,7 @@ mod decl { }: Self::Args, vm: &VirtualMachine, ) -> PyResult { - PyItertoolsFilterFalse { + Self { predicate, iterable, } @@ -1117,7 +1121,7 @@ mod decl { type Args = AccumulateArgs; fn py_new(cls: PyTypeRef, args: AccumulateArgs, vm: &VirtualMachine) -> PyResult { - PyItertoolsAccumulate { + Self { iterable: args.iterable, bin_op: args.func.flatten(), initial: args.initial.flatten(), @@ -1226,8 +1230,8 @@ mod decl { } impl PyItertoolsTeeData { - fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyResult> { - Ok(PyRc::new(PyItertoolsTeeData { + fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyResult> { + Ok(PyRc::new(Self { iterable, values: PyRwLock::new(vec![]), })) @@ -1295,7 +1299,7 @@ mod decl { if iterator.class().is(PyItertoolsTee::class(&vm.ctx)) { return vm.call_special_method(&iterator, identifier!(vm, __copy__), ()); } - Ok(PyItertoolsTee { + Ok(Self { tee_data: PyItertoolsTeeData::new(iterator, vm)?, index: AtomicCell::new(0), } @@ -1354,7 +1358,7 @@ mod decl { let l = pools.len(); - PyItertoolsProduct { + Self { pools, idxs: PyRwLock::new(vec![0; l]), cur: AtomicCell::new(l.wrapping_sub(1)), @@ -1394,8 +1398,7 @@ mod decl { fn __setstate__(zelf: PyRef, state: PyTupleRef, vm: &VirtualMachine) -> PyResult<()> { let args = state.as_slice(); if args.len() != zelf.pools.len() { - let msg = "Invalid number of arguments".to_string(); - return Err(vm.new_type_error(msg)); + return Err(vm.new_type_error("Invalid number of arguments")); } let mut idxs: PyRwLockWriteGuard<'_, Vec> = zelf.idxs.write(); idxs.clear(); @@ -1642,7 +1645,7 @@ mod decl { let n = pool.len(); - PyItertoolsCombinationsWithReplacement { + Self { pool, indices: PyRwLock::new(vec![0; r]), r: AtomicCell::new(r), @@ -1860,7 +1863,7 @@ mod decl { fn py_new(cls: PyTypeRef, (iterators, args): Self::Args, vm: &VirtualMachine) -> PyResult { let fillvalue = args.fillvalue.unwrap_or_none(vm); let iterators = iterators.into_vec(); - PyItertoolsZipLongest { + Self { iterators, fillvalue: PyRwLock::new(fillvalue), } @@ -1943,7 +1946,7 @@ mod decl { type Args = PyIter; fn py_new(cls: PyTypeRef, iterator: Self::Args, vm: &VirtualMachine) -> PyResult { - PyItertoolsPairwise { + Self { iterator, old: PyRwLock::new(None), } From e2ea8fcedac8d227fc274b0ec3ba60ae4f7de24c Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 12:31:12 +0300 Subject: [PATCH 2/9] Add macro --- vm/src/stdlib/itertools.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index e511548ade..80b48f559f 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -1,5 +1,14 @@ pub(crate) use decl::make_module; +macro_rules! handle_pyiter_return { + ($input:expr) => { + match $input { + PyIterReturn::Return(obj) => obj, + PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), + } + }; +} + #[pymodule(name = "itertools")] mod decl { use crate::stdlib::itertools::decl::int::get_value; From 1f93b9578787b309ade5d6a651715af7a9a89e19 Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 12:43:25 +0300 Subject: [PATCH 3/9] Eeplace duplicated code with macro --- vm/src/stdlib/itertools.rs | 85 +++++++------------------------------- 1 file changed, 15 insertions(+), 70 deletions(-) diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 80b48f559f..e1e82d14b1 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -225,10 +225,7 @@ mod decl { impl IterNext for PyItertoolsCompress { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = match zelf.selectors.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let sel_obj = handle_pyiter_return!(zelf.selectors.next(vm)?); let verdict = sel_obj.clone().try_to_bool(vm)?; let data_obj = zelf.data.next(vm)?; @@ -579,10 +576,7 @@ mod decl { } // might be StopIteration or anything else, which is propagated upwards - let obj = match zelf.iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let obj = handle_pyiter_return!(zelf.iterable.next(vm)?); let predicate = &zelf.predicate; let verdict = predicate.call((obj.clone(),), vm)?; @@ -667,12 +661,7 @@ mod decl { if !zelf.start_flag.load() { loop { - let obj = match iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - return Ok(PyIterReturn::StopIteration(v)); - } - }; + let obj = handle_pyiter_return!(iterable.next(vm)?); let pred = predicate.clone(); let pred_value = pred.invoke((obj.clone(),), vm)?; if !pred_value.try_to_bool(vm)? { @@ -766,10 +755,7 @@ mod decl { &self, vm: &VirtualMachine, ) -> PyResult> { - let new_value = match self.iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let new_value = handle_pyiter_return!(self.iterable.next(vm)?); let new_key = if let Some(ref kf) = self.key_func { kf.call((new_value.clone(),), vm)? } else { @@ -793,23 +779,13 @@ mod decl { let (value, key) = if let Some(old_key) = current_key { loop { - let (value, new_key) = match zelf.advance(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - return Ok(PyIterReturn::StopIteration(v)); - } - }; + let (value, new_key) = handle_pyiter_return!(zelf.advance(vm)?); if !vm.bool_eq(&new_key, &old_key)? { break (value, new_key); } } } else { - match zelf.advance(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - return Ok(PyIterReturn::StopIteration(v)); - } - } + handle_pyiter_return!(zelf.advance(vm)?) }; state = zelf.state.lock(); @@ -859,10 +835,7 @@ mod decl { state.current_key.as_ref().unwrap().clone() }; - let (value, key) = match zelf.groupby.advance(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let (value, key) = handle_pyiter_return!(zelf.groupby.advance(vm)?); if vm.bool_eq(&key, &old_key)? { Ok(PyIterReturn::Return(value)) } else { @@ -1021,10 +994,7 @@ mod decl { } } - let obj = match zelf.iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let obj = handle_pyiter_return!(zelf.iterable.next(vm)?); zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. @@ -1090,10 +1060,7 @@ mod decl { let iterable = &zelf.iterable; loop { - let obj = match iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let obj = handle_pyiter_return!(iterable.next(vm)?); let pred_value = if vm.is_none(predicate) { obj.clone() } else { @@ -1205,21 +1172,11 @@ mod decl { let next_acc_value = match acc_value { None => match &zelf.initial { - None => match iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - return Ok(PyIterReturn::StopIteration(v)); - } - }, + None => handle_pyiter_return!(iterable.next(vm)?), Some(obj) => obj.clone(), }, Some(value) => { - let obj = match iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => { - return Ok(PyIterReturn::StopIteration(v)); - } - }; + let obj = handle_pyiter_return!(iterable.next(vm)?); match &zelf.bin_op { None => vm._add(&value, &obj)?, Some(op) => op.call((value, obj), vm)?, @@ -1248,10 +1205,7 @@ mod decl { fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { if self.values.read().len() == index { - let result = match self.iterable.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let result = handle_pyiter_return!(self.iterable.next(vm)?); self.values.write().push(result); } Ok(PyIterReturn::Return(self.values.read()[index].clone())) @@ -1327,10 +1281,7 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let value = match zelf.tee_data.get_item(vm, zelf.index.load())? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let value = handle_pyiter_return!(zelf.tee_data.get_item(vm, zelf.index.load())?); zelf.index.fetch_add(1); Ok(PyIterReturn::Return(value)) } @@ -1972,16 +1923,10 @@ mod decl { impl IterNext for PyItertoolsPairwise { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { let old = match zelf.old.read().clone() { - None => match zelf.iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }, + None => handle_pyiter_return!(zelf.iterator.next(vm)?), Some(obj) => obj, }; - let new = match zelf.iterator.next(vm)? { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - }; + let new = handle_pyiter_return!(zelf.iterator.next(vm)?); *zelf.old.write() = Some(new.clone()); Ok(PyIterReturn::Return(vm.new_tuple((old, new)).into())) } From 9184d7cc0778dfb232711d16b7f1711a1db4daca Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 15:28:49 +0300 Subject: [PATCH 4/9] Move macro to protocol/iter.rs --- vm/src/protocol/iter.rs | 12 ++++++++++++ vm/src/stdlib/itertools.rs | 11 +---------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index 34f1b01ca8..ae38080d12 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -276,3 +276,15 @@ where (self.length_hint.unwrap_or(0), self.length_hint) } } + +#[macro_export] +macro_rules! handle_pyiter_return { + ($input:expr) => { + match $input { + $crate::protocol::PyIterReturn::Return(obj) => obj, + $crate::protocol::PyIterReturn::StopIteration(v) => { + return Ok($crate::protocol::PyIterReturn::StopIteration(v)) + } + } + }; +} diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index ba502362dc..3cae370ac6 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -1,14 +1,5 @@ pub(crate) use decl::make_module; -macro_rules! handle_pyiter_return { - ($input:expr) => { - match $input { - PyIterReturn::Return(obj) => obj, - PyIterReturn::StopIteration(v) => return Ok(PyIterReturn::StopIteration(v)), - } - }; -} - #[pymodule(name = "itertools")] mod decl { use crate::stdlib::itertools::decl::int::get_value; @@ -25,7 +16,7 @@ mod decl { }, convert::ToPyObject, function::{ArgCallable, ArgIntoBool, FuncArgs, OptionalArg, OptionalOption, PosArgs}, - identifier, + handle_pyiter_return, identifier, protocol::{PyIter, PyIterReturn, PyNumber}, stdlib::sys, types::{Constructor, IterNext, Iterable, Representable, SelfIter}, From 1a8a1aba0a23750489172246d98d6567247be69d Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 15:31:03 +0300 Subject: [PATCH 5/9] Rename macro --- vm/src/protocol/iter.rs | 2 +- vm/src/stdlib/itertools.rs | 31 ++++++++++++++++--------------- 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index ae38080d12..a887702851 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -278,7 +278,7 @@ where } #[macro_export] -macro_rules! handle_pyiter_return { +macro_rules! raise_stop { ($input:expr) => { match $input { $crate::protocol::PyIterReturn::Return(obj) => obj, diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 3cae370ac6..16d0c6b722 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -16,8 +16,9 @@ mod decl { }, convert::ToPyObject, function::{ArgCallable, ArgIntoBool, FuncArgs, OptionalArg, OptionalOption, PosArgs}, - handle_pyiter_return, identifier, + identifier, protocol::{PyIter, PyIterReturn, PyNumber}, + raise_stop, stdlib::sys, types::{Constructor, IterNext, Iterable, Representable, SelfIter}, }; @@ -216,7 +217,7 @@ mod decl { impl IterNext for PyItertoolsCompress { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = handle_pyiter_return!(zelf.selectors.next(vm)?); + let sel_obj = raise_stop!(zelf.selectors.next(vm)?); let verdict = sel_obj.clone().try_to_bool(vm)?; let data_obj = zelf.data.next(vm)?; @@ -567,7 +568,7 @@ mod decl { } // might be StopIteration or anything else, which is propagated upwards - let obj = handle_pyiter_return!(zelf.iterable.next(vm)?); + let obj = raise_stop!(zelf.iterable.next(vm)?); let predicate = &zelf.predicate; let verdict = predicate.call((obj.clone(),), vm)?; @@ -652,7 +653,7 @@ mod decl { if !zelf.start_flag.load() { loop { - let obj = handle_pyiter_return!(iterable.next(vm)?); + let obj = raise_stop!(iterable.next(vm)?); let pred = predicate.clone(); let pred_value = pred.invoke((obj.clone(),), vm)?; if !pred_value.try_to_bool(vm)? { @@ -746,7 +747,7 @@ mod decl { &self, vm: &VirtualMachine, ) -> PyResult> { - let new_value = handle_pyiter_return!(self.iterable.next(vm)?); + let new_value = raise_stop!(self.iterable.next(vm)?); let new_key = if let Some(ref kf) = self.key_func { kf.call((new_value.clone(),), vm)? } else { @@ -770,13 +771,13 @@ mod decl { let (value, key) = if let Some(old_key) = current_key { loop { - let (value, new_key) = handle_pyiter_return!(zelf.advance(vm)?); + let (value, new_key) = raise_stop!(zelf.advance(vm)?); if !vm.bool_eq(&new_key, &old_key)? { break (value, new_key); } } } else { - handle_pyiter_return!(zelf.advance(vm)?) + raise_stop!(zelf.advance(vm)?) }; state = zelf.state.lock(); @@ -826,7 +827,7 @@ mod decl { state.current_key.as_ref().unwrap().clone() }; - let (value, key) = handle_pyiter_return!(zelf.groupby.advance(vm)?); + let (value, key) = raise_stop!(zelf.groupby.advance(vm)?); if vm.bool_eq(&key, &old_key)? { Ok(PyIterReturn::Return(value)) } else { @@ -985,7 +986,7 @@ mod decl { } } - let obj = handle_pyiter_return!(zelf.iterable.next(vm)?); + let obj = raise_stop!(zelf.iterable.next(vm)?); zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. @@ -1051,7 +1052,7 @@ mod decl { let iterable = &zelf.iterable; loop { - let obj = handle_pyiter_return!(iterable.next(vm)?); + let obj = raise_stop!(iterable.next(vm)?); let pred_value = if vm.is_none(predicate) { obj.clone() } else { @@ -1163,11 +1164,11 @@ mod decl { let next_acc_value = match acc_value { None => match &zelf.initial { - None => handle_pyiter_return!(iterable.next(vm)?), + None => raise_stop!(iterable.next(vm)?), Some(obj) => obj.clone(), }, Some(value) => { - let obj = handle_pyiter_return!(iterable.next(vm)?); + let obj = raise_stop!(iterable.next(vm)?); match &zelf.bin_op { None => vm._add(&value, &obj)?, Some(op) => op.call((value, obj), vm)?, @@ -1196,7 +1197,7 @@ mod decl { fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { if self.values.read().len() == index { - let result = handle_pyiter_return!(self.iterable.next(vm)?); + let result = raise_stop!(self.iterable.next(vm)?); self.values.write().push(result); } Ok(PyIterReturn::Return(self.values.read()[index].clone())) @@ -1272,7 +1273,7 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let value = handle_pyiter_return!(zelf.tee_data.get_item(vm, zelf.index.load())?); + let value = raise_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?); zelf.index.fetch_add(1); Ok(PyIterReturn::Return(value)) } @@ -1929,7 +1930,7 @@ mod decl { Some(obj) => obj, }; - let new = handle_pyiter_return!(zelf.iterator.next(vm)?); + let new = raise_stop!(zelf.iterator.next(vm)?); *zelf.old.write() = Some(new.clone()); Ok(PyIterReturn::Return(vm.new_tuple((old, new)).into())) From 84167d2191c8f7cd53e72199f97dd1fe6456e4eb Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 15:35:54 +0300 Subject: [PATCH 6/9] Add docs suggestion by coderabbit --- vm/src/protocol/iter.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index a887702851..79639958c0 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -277,6 +277,19 @@ where } } +/// Macro to handle `PyIterReturn` values in iterator implementations. +/// +/// Extracts the object from `PyIterReturn::Return(obj)` or performs early return +/// for `PyIterReturn::StopIteration(v)`. This macro should only be used within +/// functions that return `PyResult`. +/// +/// # Example +/// ```rust +/// fn iterator_next(&self, vm: &VirtualMachine) -> PyResult { +/// let value = raise_stop!(some_pyiter_return_value); +/// // Process the extracted value... +/// } +/// ``` #[macro_export] macro_rules! raise_stop { ($input:expr) => { From 0282fd3f424d11ce9a39dfed79da22988899de2a Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 15:46:16 +0300 Subject: [PATCH 7/9] Don't set example --- vm/src/protocol/iter.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index 79639958c0..5190d7289d 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -282,14 +282,6 @@ where /// Extracts the object from `PyIterReturn::Return(obj)` or performs early return /// for `PyIterReturn::StopIteration(v)`. This macro should only be used within /// functions that return `PyResult`. -/// -/// # Example -/// ```rust -/// fn iterator_next(&self, vm: &VirtualMachine) -> PyResult { -/// let value = raise_stop!(some_pyiter_return_value); -/// // Process the extracted value... -/// } -/// ``` #[macro_export] macro_rules! raise_stop { ($input:expr) => { From a8d3dbeb760ce32d1aca6dbeeacbbcb171e42ffe Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Wed, 2 Jul 2025 23:01:58 +0900 Subject: [PATCH 8/9] rename to raise_if_stop --- vm/src/protocol/iter.rs | 2 +- vm/src/stdlib/itertools.rs | 30 +++++++++++++++--------------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index 5190d7289d..d05bd26a8a 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -283,7 +283,7 @@ where /// for `PyIterReturn::StopIteration(v)`. This macro should only be used within /// functions that return `PyResult`. #[macro_export] -macro_rules! raise_stop { +macro_rules! raise_if_stop { ($input:expr) => { match $input { $crate::protocol::PyIterReturn::Return(obj) => obj, diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 16d0c6b722..e8f5d21c54 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -18,7 +18,7 @@ mod decl { function::{ArgCallable, ArgIntoBool, FuncArgs, OptionalArg, OptionalOption, PosArgs}, identifier, protocol::{PyIter, PyIterReturn, PyNumber}, - raise_stop, + raise_if_stop, stdlib::sys, types::{Constructor, IterNext, Iterable, Representable, SelfIter}, }; @@ -217,7 +217,7 @@ mod decl { impl IterNext for PyItertoolsCompress { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = raise_stop!(zelf.selectors.next(vm)?); + let sel_obj = raise_if_stop!(zelf.selectors.next(vm)?); let verdict = sel_obj.clone().try_to_bool(vm)?; let data_obj = zelf.data.next(vm)?; @@ -568,7 +568,7 @@ mod decl { } // might be StopIteration or anything else, which is propagated upwards - let obj = raise_stop!(zelf.iterable.next(vm)?); + let obj = raise_if_stop!(zelf.iterable.next(vm)?); let predicate = &zelf.predicate; let verdict = predicate.call((obj.clone(),), vm)?; @@ -653,7 +653,7 @@ mod decl { if !zelf.start_flag.load() { loop { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); let pred = predicate.clone(); let pred_value = pred.invoke((obj.clone(),), vm)?; if !pred_value.try_to_bool(vm)? { @@ -747,7 +747,7 @@ mod decl { &self, vm: &VirtualMachine, ) -> PyResult> { - let new_value = raise_stop!(self.iterable.next(vm)?); + let new_value = raise_if_stop!(self.iterable.next(vm)?); let new_key = if let Some(ref kf) = self.key_func { kf.call((new_value.clone(),), vm)? } else { @@ -771,13 +771,13 @@ mod decl { let (value, key) = if let Some(old_key) = current_key { loop { - let (value, new_key) = raise_stop!(zelf.advance(vm)?); + let (value, new_key) = raise_if_stop!(zelf.advance(vm)?); if !vm.bool_eq(&new_key, &old_key)? { break (value, new_key); } } } else { - raise_stop!(zelf.advance(vm)?) + raise_if_stop!(zelf.advance(vm)?) }; state = zelf.state.lock(); @@ -827,7 +827,7 @@ mod decl { state.current_key.as_ref().unwrap().clone() }; - let (value, key) = raise_stop!(zelf.groupby.advance(vm)?); + let (value, key) = raise_if_stop!(zelf.groupby.advance(vm)?); if vm.bool_eq(&key, &old_key)? { Ok(PyIterReturn::Return(value)) } else { @@ -986,7 +986,7 @@ mod decl { } } - let obj = raise_stop!(zelf.iterable.next(vm)?); + let obj = raise_if_stop!(zelf.iterable.next(vm)?); zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. @@ -1052,7 +1052,7 @@ mod decl { let iterable = &zelf.iterable; loop { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); let pred_value = if vm.is_none(predicate) { obj.clone() } else { @@ -1164,11 +1164,11 @@ mod decl { let next_acc_value = match acc_value { None => match &zelf.initial { - None => raise_stop!(iterable.next(vm)?), + None => raise_if_stop!(iterable.next(vm)?), Some(obj) => obj.clone(), }, Some(value) => { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); match &zelf.bin_op { None => vm._add(&value, &obj)?, Some(op) => op.call((value, obj), vm)?, @@ -1197,7 +1197,7 @@ mod decl { fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { if self.values.read().len() == index { - let result = raise_stop!(self.iterable.next(vm)?); + let result = raise_if_stop!(self.iterable.next(vm)?); self.values.write().push(result); } Ok(PyIterReturn::Return(self.values.read()[index].clone())) @@ -1273,7 +1273,7 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let value = raise_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?); + let value = raise_if_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?); zelf.index.fetch_add(1); Ok(PyIterReturn::Return(value)) } @@ -1930,7 +1930,7 @@ mod decl { Some(obj) => obj, }; - let new = raise_stop!(zelf.iterator.next(vm)?); + let new = raise_if_stop!(zelf.iterator.next(vm)?); *zelf.old.write() = Some(new.clone()); Ok(PyIterReturn::Return(vm.new_tuple((old, new)).into())) From 1eeafe06cac6823a600d2225e26e54bd20810f98 Mon Sep 17 00:00:00 2001 From: ShaharNaveh <50263213+ShaharNaveh@users.noreply.github.com> Date: Wed, 2 Jul 2025 18:16:29 +0300 Subject: [PATCH 9/9] Rename to `raise_if_stop` --- vm/src/protocol/iter.rs | 2 +- vm/src/stdlib/itertools.rs | 30 +++++++++++++++--------------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/vm/src/protocol/iter.rs b/vm/src/protocol/iter.rs index 5190d7289d..d05bd26a8a 100644 --- a/vm/src/protocol/iter.rs +++ b/vm/src/protocol/iter.rs @@ -283,7 +283,7 @@ where /// for `PyIterReturn::StopIteration(v)`. This macro should only be used within /// functions that return `PyResult`. #[macro_export] -macro_rules! raise_stop { +macro_rules! raise_if_stop { ($input:expr) => { match $input { $crate::protocol::PyIterReturn::Return(obj) => obj, diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index 16d0c6b722..e8f5d21c54 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -18,7 +18,7 @@ mod decl { function::{ArgCallable, ArgIntoBool, FuncArgs, OptionalArg, OptionalOption, PosArgs}, identifier, protocol::{PyIter, PyIterReturn, PyNumber}, - raise_stop, + raise_if_stop, stdlib::sys, types::{Constructor, IterNext, Iterable, Representable, SelfIter}, }; @@ -217,7 +217,7 @@ mod decl { impl IterNext for PyItertoolsCompress { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { loop { - let sel_obj = raise_stop!(zelf.selectors.next(vm)?); + let sel_obj = raise_if_stop!(zelf.selectors.next(vm)?); let verdict = sel_obj.clone().try_to_bool(vm)?; let data_obj = zelf.data.next(vm)?; @@ -568,7 +568,7 @@ mod decl { } // might be StopIteration or anything else, which is propagated upwards - let obj = raise_stop!(zelf.iterable.next(vm)?); + let obj = raise_if_stop!(zelf.iterable.next(vm)?); let predicate = &zelf.predicate; let verdict = predicate.call((obj.clone(),), vm)?; @@ -653,7 +653,7 @@ mod decl { if !zelf.start_flag.load() { loop { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); let pred = predicate.clone(); let pred_value = pred.invoke((obj.clone(),), vm)?; if !pred_value.try_to_bool(vm)? { @@ -747,7 +747,7 @@ mod decl { &self, vm: &VirtualMachine, ) -> PyResult> { - let new_value = raise_stop!(self.iterable.next(vm)?); + let new_value = raise_if_stop!(self.iterable.next(vm)?); let new_key = if let Some(ref kf) = self.key_func { kf.call((new_value.clone(),), vm)? } else { @@ -771,13 +771,13 @@ mod decl { let (value, key) = if let Some(old_key) = current_key { loop { - let (value, new_key) = raise_stop!(zelf.advance(vm)?); + let (value, new_key) = raise_if_stop!(zelf.advance(vm)?); if !vm.bool_eq(&new_key, &old_key)? { break (value, new_key); } } } else { - raise_stop!(zelf.advance(vm)?) + raise_if_stop!(zelf.advance(vm)?) }; state = zelf.state.lock(); @@ -827,7 +827,7 @@ mod decl { state.current_key.as_ref().unwrap().clone() }; - let (value, key) = raise_stop!(zelf.groupby.advance(vm)?); + let (value, key) = raise_if_stop!(zelf.groupby.advance(vm)?); if vm.bool_eq(&key, &old_key)? { Ok(PyIterReturn::Return(value)) } else { @@ -986,7 +986,7 @@ mod decl { } } - let obj = raise_stop!(zelf.iterable.next(vm)?); + let obj = raise_if_stop!(zelf.iterable.next(vm)?); zelf.cur.fetch_add(1); // TODO is this overflow check required? attempts to copy CPython. @@ -1052,7 +1052,7 @@ mod decl { let iterable = &zelf.iterable; loop { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); let pred_value = if vm.is_none(predicate) { obj.clone() } else { @@ -1164,11 +1164,11 @@ mod decl { let next_acc_value = match acc_value { None => match &zelf.initial { - None => raise_stop!(iterable.next(vm)?), + None => raise_if_stop!(iterable.next(vm)?), Some(obj) => obj.clone(), }, Some(value) => { - let obj = raise_stop!(iterable.next(vm)?); + let obj = raise_if_stop!(iterable.next(vm)?); match &zelf.bin_op { None => vm._add(&value, &obj)?, Some(op) => op.call((value, obj), vm)?, @@ -1197,7 +1197,7 @@ mod decl { fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { if self.values.read().len() == index { - let result = raise_stop!(self.iterable.next(vm)?); + let result = raise_if_stop!(self.iterable.next(vm)?); self.values.write().push(result); } Ok(PyIterReturn::Return(self.values.read()[index].clone())) @@ -1273,7 +1273,7 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let value = raise_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?); + let value = raise_if_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?); zelf.index.fetch_add(1); Ok(PyIterReturn::Return(value)) } @@ -1930,7 +1930,7 @@ mod decl { Some(obj) => obj, }; - let new = raise_stop!(zelf.iterator.next(vm)?); + let new = raise_if_stop!(zelf.iterator.next(vm)?); *zelf.old.write() = Some(new.clone()); Ok(PyIterReturn::Return(vm.new_tuple((old, new)).into()))