ArgIndex for bytes/bytearray

vm/src/builtins/bytearray.rs

@@ -20,9 +20,9 @@ use crate::{
         },
     },
     convert::{ToPyObject, ToPyResult},
-    function::Either,
     function::{
-        ArgBytesLike, ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue,
+        ArgBytesLike, ArgIterable, ArgSize, Either, FuncArgs, OptionalArg, OptionalOption,
+        PyComparisonValue,
     },
     protocol::{
         BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
@@ -67,6 +67,10 @@ impl PyByteArray {
     pub fn borrow_buf_mut(&self) -> PyMappedRwLockWriteGuard<'_, Vec<u8>> {
         PyRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements)
     }
+
+    fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
+        self.inner().mul(value, vm).map(|x| x.into())
+    }
 }
 
 impl From<PyBytesInner> for PyByteArray {
@@ -658,13 +662,13 @@ impl PyByteArray {
 
     #[pymethod(name = "__rmul__")]
     #[pymethod(magic)]
-    fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
-        self.inner().mul(value, vm).map(|x| x.into())
+    fn mul(&self, value: ArgSize, vm: &VirtualMachine) -> PyResult<Self> {
+        self.repeat(value.into(), vm)
     }
 
     #[pymethod(magic)]
-    fn imul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-        Self::irepeat(&zelf, value, vm)?;
+    fn imul(zelf: PyRef<Self>, value: ArgSize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        Self::irepeat(&zelf, value.into(), vm)?;
         Ok(zelf)
     }
 
@@ -820,7 +824,7 @@ impl AsSequence for PyByteArray {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 PyByteArray::sequence_downcast(seq)
-                    .mul(n, vm)
+                    .repeat(n, vm)
                     .map(|x| x.into_pyobject(vm))
             }),
             item: atomic_func!(|seq, i, vm| {
@@ -849,7 +853,8 @@ impl AsSequence for PyByteArray {
             }),
             inplace_repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyByteArray::sequence_downcast(seq).to_owned();
-                PyByteArray::imul(zelf, n, vm).map(|x| x.into())
+                PyByteArray::irepeat(&zelf, n, vm)?;
+                Ok(zelf.into())
             }),
         };
         &AS_SEQUENCE

vm/src/builtins/bytes.rs

@@ -11,8 +11,9 @@ use crate::{
     class::PyClassImpl,
     common::{hash::PyHash, lock::PyMutex},
     convert::{ToPyObject, ToPyResult},
-    function::Either,
-    function::{ArgBytesLike, ArgIterable, OptionalArg, OptionalOption, PyComparisonValue},
+    function::{
+        ArgBytesLike, ArgIndex, ArgIterable, Either, OptionalArg, OptionalOption, PyComparisonValue,
+    },
     protocol::{
         BufferDescriptor, BufferMethods, PyBuffer, PyIterReturn, PyMappingMethods, PyNumberMethods,
         PySequenceMethods,
@@ -99,6 +100,19 @@ impl PyBytes {
     pub fn new_ref(data: Vec<u8>, ctx: &Context) -> PyRef<Self> {
         PyRef::new_ref(Self::from(data), ctx.types.bytes_type.to_owned(), None)
     }
+
+    fn repeat(zelf: PyRef<Self>, count: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        if count == 1 && zelf.class().is(vm.ctx.types.bytes_type) {
+            // Special case: when some `bytes` is multiplied by `1`,
+            // nothing really happens, we need to return an object itself
+            // with the same `id()` to be compatible with CPython.
+            // This only works for `bytes` itself, not its subclasses.
+            return Ok(zelf);
+        }
+        zelf.inner
+            .mul(count, vm)
+            .map(|x| Self::from(x).into_ref(vm))
+    }
 }
 
 #[pyclass(
@@ -497,17 +511,8 @@ impl PyBytes {
 
     #[pymethod(name = "__rmul__")]
     #[pymethod(magic)]
-    fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
-        if value == 1 && zelf.class().is(vm.ctx.types.bytes_type) {
-            // Special case: when some `bytes` is multiplied by `1`,
-            // nothing really happens, we need to return an object itself
-            // with the same `id()` to be compatible with CPython.
-            // This only works for `bytes` itself, not its subclasses.
-            return Ok(zelf);
-        }
-        zelf.inner
-            .mul(value, vm)
-            .map(|x| Self::from(x).into_ref(vm))
+    fn mul(zelf: PyRef<Self>, value: ArgIndex, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
+        Self::repeat(zelf, value.try_to_primitive(vm)?, vm)
     }
 
     #[pymethod(name = "__mod__")]
@@ -605,7 +610,7 @@ impl AsSequence for PyBytes {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 if let Ok(zelf) = seq.obj.to_owned().downcast::<PyBytes>() {
-                    PyBytes::mul(zelf, n, vm).to_pyresult(vm)
+                    PyBytes::repeat(zelf, n, vm).to_pyresult(vm)
                 } else {
                     Err(vm.new_type_error("bad argument type for built-in operation".to_owned()))
                 }