Fix str.[r]split

coolreader18 · coolreader18 · commit dacd119d4aa7 · 2019-08-23T17:34:56.000-05:00
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -349,43 +349,81 @@ impl PyString {
     fn split(
         &self,
         pattern: OptionalArg<PyStringRef>,
-        num: OptionalArg<usize>,
+        num: OptionalArg<isize>,
         vm: &VirtualMachine,
     ) -> PyObjectRef {
         let value = &self.value;
         let pattern = match pattern {
-            OptionalArg::Present(ref s) => &s.value,
-            OptionalArg::Missing => " ",
+            OptionalArg::Present(ref s) => Some(s.as_str()),
+            OptionalArg::Missing => None,
+        };
+        let num_splits = num.into_option().unwrap_or(-1);
+        let elements = if let Some(pattern) = pattern {
+            if num_splits.is_negative() {
+                value
+                    .split(pattern)
+                    .map(|o| vm.ctx.new_str(o.to_string()))
+                    .collect()
+            } else {
+                value
+                    .splitn(num_splits as usize + 1, pattern)
+                    .map(|o| vm.ctx.new_str(o.to_string()))
+                    .collect()
+            }
+        } else if num_splits.is_negative() {
+            value
+                .split(|c: char| c.is_ascii_whitespace())
+                .filter(|s| !s.is_empty())
+                .map(|o| vm.ctx.new_str(o.to_string()))
+                .collect()
+        } else {
+            value
+                .splitn(num_splits as usize + 1, |c: char| c.is_ascii_whitespace())
+                .filter(|s| !s.is_empty())
+                .map(|o| vm.ctx.new_str(o.to_string()))
+                .collect()
         };
-        let num_splits = num
-            .into_option()
-            .unwrap_or_else(|| value.split(pattern).count());
-        let elements = value
-            .splitn(num_splits + 1, pattern)
-            .map(|o| vm.ctx.new_str(o.to_string()))
-            .collect();
         vm.ctx.new_list(elements)
     }
 
     #[pymethod]
     fn rsplit(
         &self,
         pattern: OptionalArg<PyStringRef>,
-        num: OptionalArg<usize>,
+        num: OptionalArg<isize>,
         vm: &VirtualMachine,
     ) -> PyObjectRef {
         let value = &self.value;
         let pattern = match pattern {
-            OptionalArg::Present(ref s) => &s.value,
-            OptionalArg::Missing => " ",
+            OptionalArg::Present(ref s) => Some(s.as_str()),
+            OptionalArg::Missing => None,
+        };
+        let num_splits = num.into_option().unwrap_or(-1);
+        let mut elements: Vec<_> = if let Some(pattern) = pattern {
+            if num_splits.is_negative() {
+                value
+                    .split(pattern)
+                    .map(|o| vm.ctx.new_str(o.to_string()))
+                    .collect()
+            } else {
+                value
+                    .splitn(num_splits as usize + 1, pattern)
+                    .map(|o| vm.ctx.new_str(o.to_string()))
+                    .collect()
+            }
+        } else if num_splits.is_negative() {
+            value
+                .split(|c: char| c.is_ascii_whitespace())
+                .filter(|s| !s.is_empty())
+                .map(|o| vm.ctx.new_str(o.to_string()))
+                .collect()
+        } else {
+            value
+                .splitn(num_splits as usize + 1, |c: char| c.is_ascii_whitespace())
+                .filter(|s| !s.is_empty())
+                .map(|o| vm.ctx.new_str(o.to_string()))
+                .collect()
         };
-        let num_splits = num
-            .into_option()
-            .unwrap_or_else(|| value.split(pattern).count());
-        let mut elements: Vec<_> = value
-            .rsplitn(num_splits + 1, pattern)
-            .map(|o| vm.ctx.new_str(o.to_string()))
-            .collect();
         // Unlike Python rsplit, Rust rsplitn returns an iterator that
         // starts from the end of the string.
         elements.reverse();
