Skip to content

Commit 7d1fae0

Browse files
committed
ArgSize for list
1 parent f190cdb commit 7d1fae0

File tree

1 file changed

+20
-11
lines changed

1 file changed

+20
-11
lines changed

vm/src/builtins/list.rs

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::common::lock::{
66
use crate::{
77
class::PyClassImpl,
88
convert::ToPyObject,
9-
function::{FuncArgs, OptionalArg, PyComparisonValue},
9+
function::{ArgSize, FuncArgs, OptionalArg, PyComparisonValue},
1010
iter::PyExactSizeIterator,
1111
protocol::{PyIterReturn, PyMappingMethods, PySequenceMethods},
1212
recursion::ReprGuard,
@@ -73,6 +73,17 @@ impl PyList {
7373
pub fn borrow_vec_mut(&self) -> PyRwLockWriteGuard<'_, Vec<PyObjectRef>> {
7474
self.elements.write()
7575
}
76+
77+
fn repeat(&self, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
78+
let elements = &*self.borrow_vec();
79+
let v = elements.mul(vm, n)?;
80+
Ok(Self::new_ref(v, &vm.ctx))
81+
}
82+
83+
fn irepeat(zelf: PyRef<Self>, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
84+
zelf.borrow_vec_mut().imul(vm, n)?;
85+
Ok(zelf)
86+
}
7687
}
7788

7889
#[derive(FromArgs, Default)]
@@ -232,16 +243,13 @@ impl PyList {
232243

233244
#[pymethod(magic)]
234245
#[pymethod(name = "__rmul__")]
235-
fn mul(&self, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
236-
let elements = &*self.borrow_vec();
237-
let v = elements.mul(vm, n)?;
238-
Ok(Self::new_ref(v, &vm.ctx))
246+
fn mul(&self, n: ArgSize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
247+
self.repeat(n.into(), vm)
239248
}
240249

241250
#[pymethod(magic)]
242-
fn imul(zelf: PyRef<Self>, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
243-
zelf.borrow_vec_mut().imul(vm, n)?;
244-
Ok(zelf)
251+
fn imul(zelf: PyRef<Self>, n: ArgSize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
252+
Self::irepeat(zelf, n.into(), vm)
245253
}
246254

247255
#[pymethod]
@@ -423,7 +431,9 @@ impl AsSequence for PyList {
423431
.map(|x| x.into())
424432
}),
425433
repeat: atomic_func!(|seq, n, vm| {
426-
PyList::sequence_downcast(seq).mul(n, vm).map(|x| x.into())
434+
PyList::sequence_downcast(seq)
435+
.repeat(n, vm)
436+
.map(|x| x.into())
427437
}),
428438
item: atomic_func!(|seq, i, vm| {
429439
PyList::sequence_downcast(seq)
@@ -448,8 +458,7 @@ impl AsSequence for PyList {
448458
}),
449459
inplace_repeat: atomic_func!(|seq, n, vm| {
450460
let zelf = PyList::sequence_downcast(seq);
451-
zelf.borrow_vec_mut().imul(vm, n)?;
452-
Ok(zelf.to_owned().into())
461+
Ok(PyList::irepeat(zelf.to_owned(), n, vm)?.into())
453462
}),
454463
};
455464
&AS_SEQUENCE

0 commit comments

Comments
 (0)