diff --git a/vm/src/stdlib/itertools.rs b/vm/src/stdlib/itertools.rs index a8e1b5258..20e266b75 100644 --- a/vm/src/stdlib/itertools.rs +++ b/vm/src/stdlib/itertools.rs @@ -17,7 +17,7 @@ mod decl { AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, PyWeakRef, VirtualMachine, }; use crossbeam_utils::atomic::AtomicCell; - use num_traits::{Signed, ToPrimitive}; + use num_traits::{Signed, ToPrimitive, One}; use std::fmt; #[pyattr] @@ -220,10 +220,13 @@ mod decl { } #[pymethod(magic)] - fn repr(&self) -> PyResult { - let cur = self.cur.read(); - - Ok(format!("count({})", cur)) + fn repr(&self, vm: &VirtualMachine) -> PyResult { + let cur = format!("{}", self.cur.read().clone().repr(vm)?); + let step = self.step.as_bigint(); + if step.is_one() { + return Ok(format!("count({})", cur)); + } + Ok(format!("count({}, {})", cur, step.to_string())) } } impl IterNextIterable for PyItertoolsCount {}