diff --git a/crates/stdlib/src/random.rs b/crates/stdlib/src/random.rs index 35e6473d8..52d45f16b 100644 --- a/crates/stdlib/src/random.rs +++ b/crates/stdlib/src/random.rs @@ -67,9 +67,16 @@ mod _random { } #[pymethod] - fn getrandbits(&self, k: isize, vm: &VirtualMachine) -> PyResult { + fn getrandbits(&self, k: PyObjectRef, vm: &VirtualMachine) -> PyResult { + let k_int = k.try_index(vm)?; + let k_bigint = k_int.as_bigint(); + if k_bigint.is_negative() { + return Err(vm.new_value_error("number of bits must be non-negative")); + } + let k: isize = k_int + .try_to_primitive(vm) + .map_err(|_| vm.new_overflow_error("getrandbits: number of bits too large"))?; match k { - ..0 => Err(vm.new_value_error("number of bits must be non-negative")), 0 => Ok(BigInt::zero()), mut k => { let mut rng = self.rng.lock(); diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index 2140f44e5..d1af433d7 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/crates/vm/src/stdlib/itertools.rs @@ -297,7 +297,7 @@ mod decl { struct PyRepeatNewArgs { object: PyObjectRef, #[pyarg(any, optional)] - times: OptionalArg, + times: OptionalArg, } impl Constructor for PyItertoolsRepeat { @@ -309,7 +309,8 @@ mod decl { vm: &VirtualMachine, ) -> PyResult { let times = match times.into_option() { - Some(int) => { + Some(obj) => { + let int = obj.try_index(vm)?; let val: isize = int.try_to_primitive(vm)?; // times always >= 0. Some(PyRwLock::new(val.to_usize().unwrap_or(0)))