Fix bytes.zfill and integrate with str.zfill

Author: youknowone

Lib/test/string_tests.py

@@ -840,7 +840,6 @@ def test_swapcase(self):
 
         self.checkraises(TypeError, 'hello', 'swapcase', 42)
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
     def test_zfill(self):
         self.checkequal('123', '123', 'zfill', 2)
         self.checkequal('123', '123', 'zfill', 3)

vm/src/obj/objbytearray.rs

@@ -457,7 +457,7 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "zfill")]
-    fn zfill(&self, width: PyIntRef) -> PyByteArray {
+    fn zfill(&self, width: isize) -> PyByteArray {
         self.borrow_value().zfill(width).into()
     }
 

vm/src/obj/objbyteinner.rs

@@ -1100,24 +1100,8 @@ impl PyByteInner {
         res
     }
 
-    pub fn zfill(&self, width: PyIntRef) -> Vec<u8> {
-        if let Some(value) = width.as_bigint().to_usize() {
-            if value < self.elements.len() {
-                return self.elements.to_vec();
-            }
-            let mut res = vec![];
-            if self.elements.starts_with(&[b'-']) {
-                res.push(b'-');
-                res.extend_from_slice(&vec![b'0'; value - self.elements.len()]);
-                res.extend_from_slice(&self.elements[1..]);
-            } else {
-                res.extend_from_slice(&vec![b'0'; value - self.elements.len()]);
-                res.extend_from_slice(&self.elements[0..]);
-            }
-            res
-        } else {
-            self.elements.to_vec()
-        }
+    pub fn zfill(&self, width: isize) -> Vec<u8> {
+        bytes_zfill(&self.elements, width.to_usize().unwrap_or(0))
     }
 
     pub fn replace(
@@ -1458,3 +1442,21 @@ impl PyBytesLike {
         }
     }
 }
+
+pub fn bytes_zfill(bytes: &[u8], width: usize) -> Vec<u8> {
+    if width <= bytes.len() {
+        bytes.to_vec()
+    } else {
+        let (sign, s) = match bytes.first() {
+            Some(_sign @ b'+') | Some(_sign @ b'-') => {
+                (unsafe { bytes.get_unchecked(..1) }, &bytes[1..])
+            }
+            _ => (&b""[..], bytes),
+        };
+        let mut filled = Vec::new();
+        filled.extend_from_slice(sign);
+        filled.extend_from_slice(&b"0".repeat(width - bytes.len()));
+        filled.extend_from_slice(s);
+        filled
+    }
+}

vm/src/obj/objbytes.rs

@@ -401,7 +401,7 @@ impl PyBytes {
     }
 
     #[pymethod(name = "zfill")]
-    fn zfill(&self, width: PyIntRef) -> PyBytes {
+    fn zfill(&self, width: isize) -> PyBytes {
         self.inner.zfill(width).into()
     }
 

vm/src/obj/objstr.rs

@@ -1058,17 +1058,13 @@ impl PyString {
     }
 
     #[pymethod]
-    fn zfill(&self, len: usize) -> String {
-        let value = &self.value;
-        if len <= value.len() {
-            value.to_owned()
-        } else {
-            let mut bytes = value.bytes();
-            let (sign, s) = match bytes.next() {
-                Some(sign @ b'+') | Some(sign @ b'-') => ((sign as char).to_string(), &value[1..]),
-                _ => ("".to_owned(), value.as_str()),
-            };
-            format!("{}{}{}", sign, "0".repeat(len - value.len()), s,)
+    fn zfill(&self, width: isize) -> String {
+        use super::objbyteinner::bytes_zfill;
+        unsafe {
+            String::from_utf8_unchecked(bytes_zfill(
+                self.value.as_bytes(),
+                width.to_usize().unwrap_or(0),
+            ))
         }
     }