Merge pull request RustPython#1852 from youknowone/bytes-strip

Fix {bytes|bytearray}.{|r|l}strip

Author: youknowone

Cargo.lock

@@ -777,7 +777,6 @@ def test_additional_rsplit(self):
             self.checkequal(['arf', 'barf'], b, 'rsplit', None)
             self.checkequal(['arf', 'barf'], b, 'rsplit', None, 2)
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
     def test_strip_whitespace(self):
         self.checkequal('hello', '   hello   ', 'strip')
         self.checkequal('hello   ', '   hello   ', 'lstrip')
@@ -795,7 +794,6 @@ def test_strip_whitespace(self):
         self.checkequal('   hello', '   hello   ', 'rstrip', None)
         self.checkequal('hello', 'hello', 'strip', None)
 
-    @unittest.skip("TODO: RUSTPYTHON test_bytes")
     def test_strip(self):
         # strip/lstrip/rstrip with str arg
         self.checkequal('hello', 'xyzzyhelloxyzzy', 'strip', 'xyz')

Lib/test/string_tests.py

@@ -72,6 +72,7 @@ result-like = "^0.2.1"
 foreign-types = "0.3"
 num_enum = "0.4"
 smallbox = "0.8"
+bstr = "0.2.12"
 
 flame = { version = "0.2", optional = true }
 flamer = { version = "0.3", optional = true }

vm/Cargo.toml

@@ -5,8 +5,8 @@ use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use super::objbyteinner::{
     ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
-    ByteInnerPosition, ByteInnerSplitOptions, ByteInnerSplitlinesOptions,
-    ByteInnerTranslateOptions, ByteOr, PyByteInner,
+    ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, ByteOr,
+    PyByteInner,
 };
 use super::objint::PyIntRef;
 use super::objiter;
@@ -15,7 +15,7 @@ use super::objstr::{PyString, PyStringRef};
 use super::objtuple::PyTupleRef;
 use super::objtype::PyClassRef;
 use crate::cformat::CFormatString;
-use crate::function::OptionalArg;
+use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objstr::do_cformat_string;
 use crate::pyobject::{
     Either, PyClassImpl, PyComparisonValue, PyContext, PyIterable, PyObjectRef, PyRef, PyResult,
@@ -373,27 +373,18 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "strip")]
-    fn strip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyByteArray> {
-        Ok(self
-            .borrow_value()
-            .strip(chars, ByteInnerPosition::All)?
-            .into())
+    fn strip(&self, chars: OptionalOption<PyByteInner>) -> PyByteArray {
+        self.borrow_value().strip(chars).into()
     }
 
     #[pymethod(name = "lstrip")]
-    fn lstrip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyByteArray> {
-        Ok(self
-            .borrow_value()
-            .strip(chars, ByteInnerPosition::Left)?
-            .into())
+    fn lstrip(&self, chars: OptionalOption<PyByteInner>) -> PyByteArray {
+        self.borrow_value().lstrip(chars).into()
     }
 
     #[pymethod(name = "rstrip")]
-    fn rstrip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyByteArray> {
-        Ok(self
-            .borrow_value()
-            .strip(chars, ByteInnerPosition::Right)?
-            .into())
+    fn rstrip(&self, chars: OptionalOption<PyByteInner>) -> PyByteArray {
+        self.borrow_value().rstrip(chars).into()
     }
 
     #[pymethod(name = "split")]

vm/src/obj/objbytearray.rs

@@ -1,9 +1,9 @@
-use std::convert::TryFrom;
-use std::ops::Range;
-
+use bstr::ByteSlice;
 use num_bigint::{BigInt, ToBigInt};
 use num_integer::Integer;
 use num_traits::{One, Signed, ToPrimitive, Zero};
+use std::convert::TryFrom;
+use std::ops::Range;
 
 use super::objbytearray::{PyByteArray, PyByteArrayRef};
 use super::objbytes::{PyBytes, PyBytesRef};
@@ -15,7 +15,7 @@ use super::objsequence::{is_valid_slice_arg, PySliceableSequence};
 use super::objslice::PySliceRef;
 use super::objstr::{self, adjust_indices, PyString, PyStringRef, StringRange};
 use super::objtuple::PyTupleRef;
-use crate::function::OptionalArg;
+use crate::function::{OptionalArg, OptionalOption};
 use crate::pyhash;
 use crate::pyobject::{
     Either, PyComparisonValue, PyIterable, PyObjectRef, PyResult, ThreadSafe, TryFromObject,
@@ -962,40 +962,37 @@ impl PyByteInner {
         Ok(res)
     }
 
-    pub fn strip(
-        &self,
-        chars: OptionalArg<PyByteInner>,
-        position: ByteInnerPosition,
-    ) -> PyResult<Vec<u8>> {
-        let is_valid_char = |c| {
-            if let OptionalArg::Present(ref bytes) = chars {
-                bytes.elements.contains(c)
-            } else {
-                c.is_ascii_whitespace()
-            }
+    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()
+    }
 
-        let mut start = 0;
-        let mut end = self.len();
-
-        if let ByteInnerPosition::Left | ByteInnerPosition::All = position {
-            for (i, c) in self.elements.iter().enumerate() {
-                if !is_valid_char(c) {
-                    start = i;
-                    break;
-                }
-            }
-        }
+    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()
+    }
 
-        if let ByteInnerPosition::Right | ByteInnerPosition::All = position {
-            for (i, c) in self.elements.iter().rev().enumerate() {
-                if !is_valid_char(c) {
-                    end = self.len() - i;
-                    break;
-                }
-            }
-        }
-        Ok(self.elements[start..end].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()
     }
 
     pub fn split(&self, options: ByteInnerSplitOptions, reverse: bool) -> PyResult<Vec<&[u8]>> {
@@ -1221,12 +1218,6 @@ pub trait ByteOr: ToPrimitive {
 
 impl ByteOr for BigInt {}
 
-pub enum ByteInnerPosition {
-    Left,
-    Right,
-    All,
-}
-
 fn split_slice<'a>(slice: &'a [u8], sep: &[u8], maxsplit: isize) -> Vec<&'a [u8]> {
     let mut splitted: Vec<&[u8]> = vec![];
     let mut prev_index = 0;

vm/src/obj/objbyteinner.rs

@@ -4,8 +4,7 @@ use std::ops::Deref;
 
 use super::objbyteinner::{
     ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
-    ByteInnerPosition, ByteInnerSplitOptions, ByteInnerSplitlinesOptions,
-    ByteInnerTranslateOptions, PyByteInner,
+    ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, PyByteInner,
 };
 use super::objint::PyIntRef;
 use super::objiter;
@@ -14,7 +13,7 @@ use super::objstr::{PyString, PyStringRef};
 use super::objtuple::PyTupleRef;
 use super::objtype::PyClassRef;
 use crate::cformat::CFormatString;
-use crate::function::OptionalArg;
+use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objstr::do_cformat_string;
 use crate::pyhash;
 use crate::pyobject::{
@@ -329,18 +328,18 @@ impl PyBytes {
     }
 
     #[pymethod(name = "strip")]
-    fn strip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyBytes> {
-        Ok(self.inner.strip(chars, ByteInnerPosition::All)?.into())
+    fn strip(&self, chars: OptionalOption<PyByteInner>) -> PyBytes {
+        self.inner.strip(chars).into()
     }
 
     #[pymethod(name = "lstrip")]
-    fn lstrip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyBytes> {
-        Ok(self.inner.strip(chars, ByteInnerPosition::Left)?.into())
+    fn lstrip(&self, chars: OptionalOption<PyByteInner>) -> PyBytes {
+        self.inner.lstrip(chars).into()
     }
 
     #[pymethod(name = "rstrip")]
-    fn rstrip(&self, chars: OptionalArg<PyByteInner>) -> PyResult<PyBytes> {
-        Ok(self.inner.strip(chars, ByteInnerPosition::Right)?.into())
+    fn rstrip(&self, chars: OptionalOption<PyByteInner>) -> PyBytes {
+        self.inner.rstrip(chars).into()
     }
 
     #[pymethod(name = "split")]