getrandbits and itertools.repeat now accept __index__ protocol.

youknowone · youknowone · commit 8676ba9128ef · 2026-02-28T22:58:56.000+09:00
diff --git 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<BigInt> {
+        fn getrandbits(&self, k: PyObjectRef, vm: &VirtualMachine) -> PyResult<BigInt> {
+            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
@@ -297,7 +297,7 @@ mod decl {
     struct PyRepeatNewArgs {
         object: PyObjectRef,
         #[pyarg(any, optional)]
-        times: OptionalArg<PyIntRef>,
+        times: OptionalArg<PyObjectRef>,
     }
 
     impl Constructor for PyItertoolsRepeat {
@@ -309,7 +309,8 @@ mod decl {
             vm: &VirtualMachine,
         ) -> PyResult<Self> {
             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)))
