diff --git a/common/src/str.rs b/common/src/str.rs index 7dbc870a3..a89c25f49 100644 --- a/common/src/str.rs +++ b/common/src/str.rs @@ -11,6 +11,24 @@ pub fn get_chars(s: &str, range: std::ops::Range) -> &str { &start[..start.len() - end.len()] } +pub fn zfill(bytes: &[u8], width: usize) -> Vec { + 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(std::iter::repeat(b'0').take(width - bytes.len())); + filled.extend_from_slice(s); + filled + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/vm/src/obj/pystr.rs b/vm/src/obj/pystr.rs index 3b37d4c9e..93ecf4de6 100644 --- a/vm/src/obj/pystr.rs +++ b/vm/src/obj/pystr.rs @@ -393,23 +393,8 @@ where } fn py_zfill(&self, width: isize) -> Vec { - let bytes = self.as_bytes(); let width = width.to_usize().unwrap_or(0); - 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(std::iter::repeat(b'0').take(width - bytes.len())); - filled.extend_from_slice(s); - filled - } + rustpython_common::str::zfill(self.as_bytes(), width) } fn py_iscase(&'s self, is_case: F, is_opposite: G) -> bool