diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs index ae39e234c..c3498947e 100644 --- a/vm/src/builtins/bytearray.rs +++ b/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> { PyRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements) } + + fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult { + self.inner().mul(value, vm).map(|x| x.into()) + } } impl From for PyByteArray { @@ -658,13 +662,13 @@ impl PyByteArray { #[pymethod(name = "__rmul__")] #[pymethod(magic)] - fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult { - self.inner().mul(value, vm).map(|x| x.into()) + fn mul(&self, value: ArgSize, vm: &VirtualMachine) -> PyResult { + self.repeat(value.into(), vm) } #[pymethod(magic)] - fn imul(zelf: PyRef, value: isize, vm: &VirtualMachine) -> PyResult> { - Self::irepeat(&zelf, value, vm)?; + fn imul(zelf: PyRef, value: ArgSize, vm: &VirtualMachine) -> PyResult> { + 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 diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs index a76ae966c..d91e0ef63 100644 --- a/vm/src/builtins/bytes.rs +++ b/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, ctx: &Context) -> PyRef { PyRef::new_ref(Self::from(data), ctx.types.bytes_type.to_owned(), None) } + + fn repeat(zelf: PyRef, count: isize, vm: &VirtualMachine) -> PyResult> { + 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, value: isize, vm: &VirtualMachine) -> PyResult> { - 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, value: ArgIndex, vm: &VirtualMachine) -> PyResult> { + 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::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())) }