diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs index 2adb94544..dffd6c95e 100644 --- a/vm/src/obj/objbyteinner.rs +++ b/vm/src/obj/objbyteinner.rs @@ -893,36 +893,33 @@ impl PyByteInner { } pub fn strip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim().to_owned(), - }; self.elements - .trim_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_with(|c| chars.contains(&(c as u8))), + |s| s.trim(), + ) + .to_vec() } pub fn lstrip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim_start().to_owned(), - }; self.elements - .trim_start_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_start_with(|c| chars.contains(&(c as u8))), + |s| s.trim_start(), + ) + .to_vec() } pub fn rstrip(&self, chars: OptionalOption) -> Vec { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.elements, - None => return self.elements.trim_end().to_owned(), - }; self.elements - .trim_end_with(|c| chars.contains(&(c as u8))) - .to_owned() + .py_strip( + chars, + |s, chars| s.trim_end_with(|c| chars.contains(&(c as u8))), + |s| s.trim_end(), + ) + .to_vec() } pub fn split( diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs index 6b1f57120..4be9a7d87 100644 --- a/vm/src/obj/objstr.rs +++ b/vm/src/obj/objstr.rs @@ -27,7 +27,7 @@ use crate::cformat::{ CNumberType, }; use crate::format::{FormatParseError, FormatPart, FormatPreconversor, FormatSpec, FormatString}; -use crate::function::{OptionalArg, PyFuncArgs}; +use crate::function::{OptionalArg, OptionalOption, PyFuncArgs}; use crate::pyhash; use crate::pyobject::{ Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyIterable, @@ -470,36 +470,35 @@ impl PyString { } #[pymethod] - fn strip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim().to_owned(), - }; - self.value.trim_matches(|c| chars.contains(c)).to_owned() - } - - #[pymethod] - fn lstrip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim_start().to_owned(), - }; + fn strip(&self, chars: OptionalOption) -> String { self.value - .trim_start_matches(|c| chars.contains(c)) + .py_strip( + chars, + |s, chars| s.trim_matches(|c| chars.contains(c)), + |s| s.trim(), + ) .to_owned() } #[pymethod] - fn rstrip(&self, chars: OptionalArg>) -> String { - let chars = chars.flat_option(); - let chars = match chars { - Some(ref chars) => &chars.value, - None => return self.value.trim_end().to_owned(), - }; + fn lstrip(&self, chars: OptionalOption) -> String { self.value - .trim_end_matches(|c| chars.contains(c)) + .py_strip( + chars, + |s, chars| s.trim_start_matches(|c| chars.contains(c)), + |s| s.trim_start(), + ) + .to_owned() + } + + #[pymethod] + fn rstrip(&self, chars: OptionalOption) -> String { + self.value + .py_strip( + chars, + |s, chars| s.trim_end_matches(|c| chars.contains(c)), + |s| s.trim_end(), + ) .to_owned() } diff --git a/vm/src/obj/pystr.rs b/vm/src/obj/pystr.rs index b78d3dbc5..7af1d9383 100644 --- a/vm/src/obj/pystr.rs +++ b/vm/src/obj/pystr.rs @@ -171,4 +171,22 @@ pub trait PyCommonString { Ok(false) } } + + fn py_strip<'a, S, FC, FD>( + &'a self, + chars: OptionalOption, + func_chars: FC, + func_default: FD, + ) -> &'a Self + where + S: PyCommonStringWrapper, + FC: Fn(&'a Self, &Self) -> &'a Self, + FD: Fn(&'a Self) -> &'a Self, + { + let chars = chars.flat_option(); + match chars { + Some(chars) => func_chars(self, chars.as_ref()), + None => func_default(self), + } + } }