8000 Add itertools.batched Support · moreal/RustPython@da53eee · GitHub
[go: up one dir, main page]

Skip to content

Commit da53eee

Browse files
Popproyouknowone
authored andcommitted
Add itertools.batched Support
1 parent e349b0f commit da53eee

File tree

1 file changed

+71
-0
lines changed

1 file changed

+71
-0
lines changed

vm/src/stdlib/itertools.rs

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1952,4 +1952,75 @@ mod decl {
19521952
Ok(PyIterReturn::Return(vm.new_tuple((old, new)).into()))
19531953
}
19541954
}
1955+
1956+
#[pyattr]
1957+
#[pyclass(name = "batched")]
1958+
#[derive(Debug, PyPayload)]
1959+
struct PyItertoolsBatched {
1960+
exhausted: AtomicCell<bool>,
1961+
iterable: PyIter,
1962+
n: AtomicCell<usize>,
1963+
}
1964+
1965+
#[derive(FromArgs)]
1966+
struct BatchedNewArgs {
1967+
#[pyarg(positional)]
1968+
iterable: PyIter,
1969+
#[pyarg(positional)]
1970+
n: PyIntRef,
1971+
}
1972+
1973+
impl Constructor for PyItertoolsBatched {
1974+
type Args = BatchedNewArgs;
1975+
1976+
fn py_new(
1977+
cls: PyTypeRef,
1978+
Self::Args { iterable, n }: Self::Args,
1979+
vm: &VirtualMachine,
1980+
) -> PyResult {
1981+
let n = n.as_bigint();
1982+
if n.is_negative() {
1983+
return Err(vm.new_value_error("n must be at least one".to_owned()));
1984+
}
1985+
let n = n.to_usize().unwrap();
1986+
1987+
Self {
1988+
iterable,
1989+
n: AtomicCell::new(n),
1990+
exhausted: AtomicCell::new(false),
1991+
}
1992+
.into_ref_with_type(vm, cls)
1993+
.map(Into::into)
1994+
}
1995+
}
1996+
1997+
#[pyclass(with(IterNext, Iterable, Constructor), flags(BASETYPE, HAS_DICT))]
1998+
impl PyItertoolsBatched {}
1999+
2000+
impl SelfIter for PyItertoolsBatched {}
2001+
2002+
impl IterNext for PyItertoolsBatched {
2003+
fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
2004+
if zelf.exhausted.load() {
2005+
return Ok(PyIterReturn::StopIteration(None));
2006+
}
2007+
let mut result: Vec<PyObjectRef> = Vec::new();
2008+
let n = zelf.n.load();
2009+
for _ in 0..n {
2010+
match zelf.iterable.next(vm)? {
2011+
PyIterReturn::Return(obj) => {
2012+
result.push(obj);
2013+
}
2014+
PyIterReturn::StopIteration(_) => {
2015+
zelf.exhausted.store(true);
2016+
break;
2017+
}
2018+
}
2019+
}
2020+
match result.len() {
2021+
0 => Ok(PyIterReturn::StopIteration(None)),
2022+
_ => Ok(PyIterReturn::Return(vm.ctx.new_tuple(result).into())),
2023+
}
2024+
}
2025+
}
19552026
}

0 commit comments

Comments
 (0)
0