add replace

Author: jgirardet

tests/snippets/bytes.py

@@ -574,3 +574,11 @@
 assert b"-42".zfill(5) == b"-0042"
 assert b"42".zfill(1) == b"42"
 assert b"42".zfill(-1) == b"42"
+
+# replace
+assert b"123456789123".replace(b"23", b"XX") == b'1XX4567891XX'
+assert b"123456789123".replace(b"23", b"XX", 1) == b'1XX456789123'
+assert b"123456789123".replace(b"23", b"XX", 0) == b"123456789123"
+assert b"123456789123".replace(b"23", b"XX", -1) == b'1XX4567891XX'
+assert b"123456789123".replace(b"23", b"") == b"14567891"
+

vm/src/obj/objbyteinner.rs

@@ -972,6 +972,64 @@ impl PyByteInner {
             self.elements.to_vec()
         }
     }
+
+    pub fn replace(
+        &self,
+        old: PyObjectRef,
+        new: PyObjectRef,
+        count: OptionalArg<PyIntRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult<Vec<u8>> {
+        let old = match try_as_bytes_like(&old) {
+            Some(value) => value,
+            None => {
+                return Err(
+                    vm.new_type_error(format!("a bytes-like object is required, not {}", old))
+                );
+            }
+        };
+
+        let new = match try_as_bytes_like(&new) {
+            Some(value) => value,
+            None => {
+                return Err(
+                    vm.new_type_error(format!("a bytes-like object is required, not {}", new))
+                );
+            }
+        };
+
+        let count = if let OptionalArg::Present(int) = count {
+            if let Some(value) = int.as_bigint().to_u32() {
+                value
+            } else {
+                self.elements.len() as u32
+            }
+        } else {
+            self.elements.len() as u32
+        };
+
+        let mut res = vec![];
+        let mut index = 0;
+        let mut done = 0;
+
+        let slice = &self.elements;
+        while index <= slice.len() - old.len() {
+            if done == count {
+                res.extend_from_slice(&slice[index..]);
+                break;
+            }
+            if &slice[index..index + old.len()] == old.as_slice() {
+                res.extend_from_slice(&new);
+                index += old.len();
+                done += 1;
+            } else {
+                res.push(slice[index]);
+                index += 1
+            }
+        }
+
+        Ok(res)
+    }
 }
 
 pub fn try_as_byte(obj: &PyObjectRef) -> Option<Vec<u8>> {

vm/src/obj/objbytes.rs

@@ -396,6 +396,17 @@ impl PyBytesRef {
     fn zfill(self, width: PyIntRef, vm: &VirtualMachine) -> PyResult {
         Ok(vm.ctx.new_bytes(self.inner.zfill(width)))
     }
+
+    #[pymethod(name = "replace")]
+    fn replace(
+        self,
+        old: PyObjectRef,
+        new: PyObjectRef,
+        count: OptionalArg<PyIntRef>,
+        vm: &VirtualMachine,
+    ) -> PyResult {
+        Ok(vm.ctx.new_bytes(self.inner.replace(old, new, count, vm)?))
+    }
 }
 
 #[derive(Debug)]