@@ -1952,4 +1952,75 @@ mod decl {
1952
1952
Ok ( PyIterReturn :: Return ( vm. new_tuple ( ( old, new) ) . into ( ) ) )
1953
1953
}
1954
1954
}
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
+ }
1955
2026
}
0 commit comments