RustPython · youknowone · Feb 21, 2023 · Feb 21, 2023 · Feb 21, 2023
diff --git a/vm/src/builtins/bytearray.rs b/vm/src/builtins/bytearray.rs
@@ -820,7 +820,7 @@ impl AsSequence for PyByteArray {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 PyByteArray::sequence_downcast(seq)
-                    .mul(n as isize, vm)
+                    .mul(n, vm)
                     .map(|x| x.into_pyobject(vm))
             }),
             item: atomic_func!(|seq, i, vm| {
@@ -849,7 +849,7 @@ impl AsSequence for PyByteArray {
             }),
             inplace_repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyByteArray::sequence_downcast(seq).to_owned();
-                PyByteArray::imul(zelf, n as isize, vm).map(|x| x.into())
+                PyByteArray::imul(zelf, n, vm).map(|x| x.into())
             }),
         };
         &AS_SEQUENCE

diff --git a/vm/src/builtins/bytes.rs b/vm/src/builtins/bytes.rs
@@ -604,10 +604,11 @@ impl AsSequence for PyBytes {
                     .map(|x| vm.ctx.new_bytes(x).into())
             }),
             repeat: atomic_func!(|seq, n, vm| {
-                Ok(vm
-                    .ctx
-                    .new_bytes(PyBytes::sequence_downcast(seq).repeat(n))
-                    .into())
+                if let Ok(zelf) = seq.obj.to_owned().downcast::<PyBytes>() {
+                    PyBytes::mul(zelf, n, vm).to_pyresult(vm)
+                } else {
+                    Err(vm.new_type_error("bad argument type for built-in operation".to_owned()))
+                }
             }),
             item: atomic_func!(|seq, i, vm| {
                 PyBytes::sequence_downcast(seq)

diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs
@@ -431,9 +431,7 @@ impl AsSequence for PyList {
                     .map(|x| x.into())
             }),
             repeat: atomic_func!(|seq, n, vm| {
-                PyList::sequence_downcast(seq)
-                    .mul(n as isize, vm)
-                    .map(|x| x.into())
+                PyList::sequence_downcast(seq).mul(n, vm).map(|x| x.into())
             }),
             item: atomic_func!(|seq, i, vm| {
                 PyList::sequence_downcast(seq)
@@ -458,7 +456,7 @@ impl AsSequence for PyList {
             }),
             inplace_repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyList::sequence_downcast(seq);
-                zelf.borrow_vec_mut().imul(vm, n as isize)?;
+                zelf.borrow_vec_mut().imul(vm, n)?;
                 Ok(zelf.to_owned().into())
             }),
         };

diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs
@@ -1291,7 +1291,7 @@ impl AsSequence for PyStr {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyStr::sequence_downcast(seq);
-                PyStr::mul(zelf.to_owned(), n as isize, vm).map(|x| x.into())
+                PyStr::mul(zelf.to_owned(), n, vm).map(|x| x.into())
             }),
             item: atomic_func!(|seq, i, vm| {
                 let zelf = PyStr::sequence_downcast(seq);

diff --git a/vm/src/builtins/tuple.rs b/vm/src/builtins/tuple.rs
@@ -353,7 +353,7 @@ impl AsSequence for PyTuple {
             }),
             repeat: atomic_func!(|seq, n, vm| {
                 let zelf = PyTuple::sequence_downcast(seq);
-                PyTuple::mul(zelf.to_owned(), n as isize, vm).map(|x| x.into())
+                PyTuple::mul(zelf.to_owned(), n, vm).map(|x| x.into())
             }),
             item: atomic_func!(|seq, i, vm| {
                 let zelf = PyTuple::sequence_downcast(seq);

diff --git a/vm/src/builtins/union.rs b/vm/src/builtins/union.rs
@@ -6,10 +6,10 @@ use crate::{
     builtins::{PyFrozenSet, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType},
     class::PyClassImpl,
     common::hash,
-    convert::ToPyObject,
+    convert::{ToPyObject, ToPyResult},
     function::PyComparisonValue,
-    protocol::PyMappingMethods,
-    types::{AsMapping, Comparable, GetAttr, Hashable, PyComparisonOp},
+    protocol::{PyMappingMethods, PyNumberMethods},
+    types::{AsMapping, AsNumber, Comparable, GetAttr, Hashable, PyComparisonOp},
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     VirtualMachine,
 };
@@ -35,7 +35,7 @@ impl PyPayload for PyUnion {
     }
 }
 
-#[pyclass(with(Hashable, Comparable, AsMapping), flags(BASETYPE))]
+#[pyclass(flags(BASETYPE), with(Hashable, Comparable, AsMapping, AsNumber))]
 impl PyUnion {
     pub fn new(args: PyTupleRef, vm: &VirtualMachine) -> Self {
         let parameters = make_parameters(&args, vm);
@@ -255,6 +255,18 @@ impl AsMapping for PyUnion {
     }
 }
 
+impl AsNumber for PyUnion {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            or: atomic_func!(|num, other, vm| {
+                PyUnion::or(num.obj.to_owned(), other.to_owned(), vm).to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl Comparable for PyUnion {
     fn cmp(
         zelf: &crate::Py<Self>,

diff --git a/vm/src/protocol/sequence.rs b/vm/src/protocol/sequence.rs
@@ -29,14 +29,14 @@ impl PyObject {
 pub struct PySequenceMethods {
     pub length: AtomicCell<Option<fn(PySequence, &VirtualMachine) -> PyResult<usize>>>,
     pub concat: AtomicCell<Option<fn(PySequence, &PyObject, &VirtualMachine) -> PyResult>>,
-    pub repeat: AtomicCell<Option<fn(PySequence, usize, &VirtualMachine) -> PyResult>>,
+    pub repeat: AtomicCell<Option<fn(PySequence, isize, &VirtualMachine) -> PyResult>>,
     pub item: AtomicCell<Option<fn(PySequence, isize, &VirtualMachine) -> PyResult>>,
     pub ass_item: AtomicCell<
         Option<fn(PySequence, isize, Option<PyObjectRef>, &VirtualMachine) -> PyResult<()>>,
     >,
     pub contains: AtomicCell<Option<fn(PySequence, &PyObject, &VirtualMachine) -> PyResult<bool>>>,
     pub inplace_concat: AtomicCell<Option<fn(PySequence, &PyObject, &VirtualMachine) -> PyResult>>,
-    pub inplace_repeat: AtomicCell<Option<fn(PySequence, usize, &VirtualMachine) -> PyResult>>,
+    pub inplace_repeat: AtomicCell<Option<fn(PySequence, isize, &VirtualMachine) -> PyResult>>,
 }
 
 impl Debug for PySequenceMethods {
@@ -130,7 +130,7 @@ impl PySequence<'_> {
         )))
     }
 
-    pub fn repeat(self, n: usize, vm: &VirtualMachine) -> PyResult {
+    pub fn repeat(self, n: isize, vm: &VirtualMachine) -> PyResult {
         if let Some(f) = self.methods.repeat.load() {
             return f(self, n, vm);
         }
@@ -168,7 +168,7 @@ impl PySequence<'_> {
         )))
     }
 
-    pub fn inplace_repeat(self, n: usize, vm: &VirtualMachine) -> PyResult {
+    pub fn inplace_repeat(self, n: isize, vm: &VirtualMachine) -> PyResult {
         if let Some(f) = self.methods.inplace_repeat.load() {
             return f(self, n, vm);
         }

diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs
@@ -525,7 +525,7 @@ mod _collections {
                 }),
                 repeat: atomic_func!(|seq, n, vm| {
                     PyDeque::sequence_downcast(seq)
-                        .mul(n as isize, vm)
+                        .mul(n, vm)
                         .map(|x| x.into_ref(vm).into())
                 }),
                 item: atomic_func!(|seq, i, vm| PyDeque::sequence_downcast(seq).getitem(i, vm)),
@@ -547,7 +547,7 @@ mod _collections {
                 }),
                 inplace_repeat: atomic_func!(|seq, n, vm| {
                     let zelf = PyDeque::sequence_downcast(seq);
-                    PyDeque::imul(zelf.to_owned(), n as isize, vm).map(|x| x.into())
+                    PyDeque::imul(zelf.to_owned(), n, vm).map(|x| x.into())
                 }),
             };