{str,bytes}.partition into pystr

youknowone · youknowone · commit f4376627e0da · 2020-07-25T22:08:09.000+09:00
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
@@ -432,7 +432,7 @@ impl PyByteArray {
     #[pymethod(name = "rpartition")]
     fn rpartition(&self, sep: PyByteInner, vm: &VirtualMachine) -> PyResult {
         let value = self.borrow_value();
-        let (front, has_mid, back) = value.rpartition(&sep, vm)?;
+        let (back, has_mid, front) = value.rpartition(&sep, vm)?;
         Ok(vm.ctx.new_tuple(vec![
             vm.ctx.new_bytearray(front.to_vec()),
             vm.ctx
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -869,37 +869,23 @@ impl PyByteInner {
         sub: &PyByteInner,
         vm: &VirtualMachine,
     ) -> PyResult<(Vec<u8>, bool, Vec<u8>)> {
-        if sub.elements.is_empty() {
-            return Err(vm.new_value_error("empty separator".to_owned()));
-        }
-
-        let mut sp = self.elements.splitn_str(2, &sub.elements);
-        let front = sp.next().unwrap().to_vec();
-        let (has_mid, back) = if let Some(back) = sp.next() {
-            (true, back.to_vec())
-        } else {
-            (false, Vec::new())
-        };
-        Ok((front, has_mid, back))
+        self.elements.py_partition(
+            &sub.elements,
+            || self.elements.splitn_str(2, &sub.elements),
+            vm,
+        )
     }
 
     pub fn rpartition(
         &self,
         sub: &PyByteInner,
         vm: &VirtualMachine,
     ) -> PyResult<(Vec<u8>, bool, Vec<u8>)> {
-        if sub.elements.is_empty() {
-            return Err(vm.new_value_error("empty separator".to_owned()));
-        }
-
-        let mut sp = self.elements.rsplitn_str(2, &sub.elements);
-        let back = sp.next().unwrap().to_vec();
-        let (has_mid, front) = if let Some(front) = sp.next() {
-            (true, front.to_vec())
-        } else {
-            (false, Vec::new())
-        };
-        Ok((front, has_mid, back))
+        self.elements.py_partition(
+            &sub.elements,
+            || self.elements.rsplitn_str(2, &sub.elements),
+            vm,
+        )
     }
 
     pub fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> Vec<u8> {
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -396,7 +396,7 @@ impl PyBytes {
     #[pymethod(name = "rpartition")]
     fn rpartition(&self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         let sub = PyByteInner::try_from_object(vm, sep.clone())?;
-        let (front, has_mid, back) = self.inner.rpartition(&sub, vm)?;
+        let (back, has_mid, front) = self.inner.rpartition(&sub, vm)?;
         Ok(vm.ctx.new_tuple(vec![
             vm.ctx.new_bytes(front),
             if has_mid {
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -813,43 +813,35 @@ impl PyString {
     }
 
     #[pymethod]
-    fn partition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyResult {
-        if sub.value.is_empty() {
-            return Err(vm.new_value_error("empty separator".to_owned()));
-        }
-        let mut sp = self.value.splitn(2, &sub.value);
-        let front = sp.next().unwrap();
-        let elems = if let Some(back) = sp.next() {
-            [front, &sub.value, back]
-        } else {
-            [front, "", ""]
-        };
-        Ok(vm.ctx.new_tuple(
-            elems
-                .iter()
-                .map(|&s| vm.ctx.new_str(s.to_owned()))
-                .collect(),
-        ))
+    fn partition(&self, sep: PyStringRef, vm: &VirtualMachine) -> PyResult {
+        let (front, has_mid, back) =
+            self.value
+                .py_partition(&sep.value, || self.value.splitn(2, &sep.value), vm)?;
+        Ok(vm.ctx.new_tuple(vec![
+            vm.ctx.new_str(front),
+            if has_mid {
+                sep.into_object()
+            } else {
+                vm.ctx.new_str("")
+            },
+            vm.ctx.new_str(back),
+        ]))
     }
 
     #[pymethod]
-    fn rpartition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyResult {
-        if sub.value.is_empty() {
-            return Err(vm.new_value_error("empty separator".to_owned()));
-        }
-        let mut sp = self.value.rsplitn(2, &sub.value);
-        let back = sp.next().unwrap();
-        let elems = if let Some(front) = sp.next() {
-            [front, &sub.value, back]
-        } else {
-            ["", "", back]
-        };
-        Ok(vm.ctx.new_tuple(
-            elems
-                .iter()
-                .map(|&s| vm.ctx.new_str(s.to_owned()))
-                .collect(),
-        ))
+    fn rpartition(&self, sep: PyStringRef, vm: &VirtualMachine) -> PyResult {
+        let (back, has_mid, front) =
+            self.value
+                .py_partition(&sep.value, || self.value.rsplitn(2, &sep.value), vm)?;
+        Ok(vm.ctx.new_tuple(vec![
+            vm.ctx.new_str(front),
+            if has_mid {
+                sep.into_object()
+            } else {
+                vm.ctx.new_str("")
+            },
+            vm.ctx.new_str(back),
+        ]))
     }
 
     /// Return `true` if the sequence is ASCII titlecase and the sequence is not
diff --git a/vm/src/obj/pystr.rs b/vm/src/obj/pystr.rs
@@ -325,6 +325,30 @@ where
         Ok(joined)
     }
 
+    fn py_partition<'a, F, S>(
+        &'a self,
+        sub: &Self,
+        split: F,
+        vm: &VirtualMachine,
+    ) -> PyResult<(Self::Container, bool, Self::Container)>
+    where
+        F: Fn() -> S,
+        S: std::iter::Iterator<Item = &'a Self>,
+    {
+        if sub.is_empty() {
+            return Err(vm.new_value_error("empty separator".to_owned()));
+        }
+
+        let mut sp = split();
+        let front = sp.next().unwrap().to_container();
+        let (has_mid, back) = if let Some(back) = sp.next() {
+            (true, back.to_container())
+        } else {
+            (false, Self::Container::new())
+        };
+        Ok((front, has_mid, back))
+    }
+
     fn py_removeprefix<FC>(
         &self,
         prefix: &Self::Container,
