Merge pull request RustPython#1876 from youknowone/share-strip

youknowone · web-flow · commit d3449ae79427 · 2020-04-26T01:44:04.000+09:00
shared string strip implemenetation
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -893,36 +893,33 @@ impl PyByteInner {
     }
 
     pub fn strip(&self, chars: OptionalOption<PyByteInner>) -> Vec<u8> {
-        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<PyByteInner>) -> Vec<u8> {
-        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<PyByteInner>) -> Vec<u8> {
-        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<F>(
diff --git 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,
@@ -474,36 +474,35 @@ impl PyString {
     }
 
     #[pymethod]
-    fn strip(&self, chars: OptionalArg<Option<PyStringRef>>) -> 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()
+    fn strip(&self, chars: OptionalOption<PyStringRef>) -> String {
+        self.value
+            .py_strip(
+                chars,
+                |s, chars| s.trim_matches(|c| chars.contains(c)),
+                |s| s.trim(),
+            )
+            .to_owned()
     }
 
     #[pymethod]
-    fn lstrip(&self, chars: OptionalArg<Option<PyStringRef>>) -> String {
-        let chars = chars.flat_option();
-        let chars = match chars {
-            Some(ref chars) => &chars.value,
-            None => return self.value.trim_start().to_owned(),
-        };
+    fn lstrip(&self, chars: OptionalOption<PyStringRef>) -> String {
         self.value
-            .trim_start_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: OptionalArg<Option<PyStringRef>>) -> String {
-        let chars = chars.flat_option();
-        let chars = match chars {
-            Some(ref chars) => &chars.value,
-            None => return self.value.trim_end().to_owned(),
-        };
+    fn rstrip(&self, chars: OptionalOption<PyStringRef>) -> String {
         self.value
-            .trim_end_matches(|c| chars.contains(c))
+            .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
@@ -171,4 +171,22 @@ pub trait PyCommonString<E> {
             Ok(false)
         }
     }
+
+    fn py_strip<'a, S, FC, FD>(
+        &'a self,
+        chars: OptionalOption<S>,
+        func_chars: FC,
+        func_default: FD,
+    ) -> &'a Self
+    where
+        S: PyCommonStringWrapper<Self>,
+        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),
+        }
+    }
 }
