Merge pull request RustPython#1859 from youknowone/strbytes-mul

palaviv · web-flow · commit 559c596bcb8f · 2020-04-11T19:22:01.000+03:00
Fix {str,bytes}.__mul__
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -421,13 +421,12 @@ impl PyBytes {
     }
 
     #[pymethod(name = "__mul__")]
-    fn repeat(&self, n: isize) -> PyBytes {
-        self.inner.repeat(n).into()
-    }
-
     #[pymethod(name = "__rmul__")]
-    fn rmul(&self, n: isize) -> PyBytes {
-        self.repeat(n)
+    fn repeat(&self, value: isize, vm: &VirtualMachine) -> PyResult<PyBytes> {
+        if value > 0 && self.inner.len() as isize > std::isize::MAX / value {
+            return Err(vm.new_overflow_error("repeated bytes are too long".to_owned()));
+        }
+        Ok(self.inner.repeat(value).into())
     }
 
     fn do_cformat(
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -332,19 +332,9 @@ impl PyString {
     }
 
     #[pymethod(name = "__mul__")]
-    fn mul(&self, multiplier: isize, vm: &VirtualMachine) -> PyResult<String> {
-        multiplier
-            .max(0)
-            .to_usize()
-            .map(|multiplier| self.value.repeat(multiplier))
-            .ok_or_else(|| {
-                vm.new_overflow_error("cannot fit 'int' into an index-sized integer".to_owned())
-            })
-    }
-
     #[pymethod(name = "__rmul__")]
-    fn rmul(&self, val: isize, vm: &VirtualMachine) -> PyResult<String> {
-        self.mul(val, vm)
+    fn mul(&self, value: isize) -> String {
+        self.value.repeat(value.max(0) as usize)
     }
 
     #[pymethod(name = "__str__")]
