use PyBytinner as arg

jgirardet · jgirardet · commit 0d505adbb513 · 2019-04-30T22:01:45.000+02:00
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -281,36 +281,35 @@ impl ByteInnerSplitOptions {
 #[derive(FromArgs)]
 pub struct ByteInnerExpandtabsOptions {
     #[pyarg(positional_or_keyword, optional = true)]
-    tabsize: OptionalArg<PyObjectRef>,
+    tabsize: OptionalArg<PyIntRef>,
 }
 
 impl ByteInnerExpandtabsOptions {
     pub fn get_value(self) -> usize {
-        if let OptionalArg::Present(value) = self.tabsize {
-            if let Some(v) = objint::get_value(&value).to_usize() {
-                v
-            } else {
-                0
-            }
-        } else {
-            8
+        match self.tabsize.into_option() {
+            Some(int) => int.as_bigint().to_usize().unwrap_or(0),
+            None => 8,
         }
     }
 }
 
 #[derive(FromArgs)]
 pub struct ByteInnerSplitlinesOptions {
     #[pyarg(positional_or_keyword, optional = true)]
-    keepends: OptionalArg<PyObjectRef>,
+    keepends: OptionalArg<bool>,
 }
 
 impl ByteInnerSplitlinesOptions {
-    pub fn get_value(self, vm: &VirtualMachine) -> PyResult<bool> {
-        if let OptionalArg::Present(value) = self.keepends {
-            Ok(bool::try_from_object(vm, value)?)
-        } else {
-            Ok(false)
+    pub fn get_value(self) -> bool {
+        match self.keepends.into_option() {
+            Some(x) => x,
+            None => false,
         }
+        // if let OptionalArg::Present(value) = self.keepends {
+        //     Ok(bool::try_from_object(vm, value)?)
+        // } else {
+        //     Ok(false)
+        // }
     }
 }
 
@@ -836,24 +835,11 @@ impl PyByteInner {
         }
     }
 
-    pub fn partition(
-        &self,
-        sep: &PyObjectRef,
-        reverse: bool,
-        vm: &VirtualMachine,
-    ) -> PyResult<(Vec<u8>, Vec<u8>)> {
-        let sep = match try_as_bytes_like(&sep) {
-            Some(value) => value,
-            None => {
-                return Err(
-                    vm.new_type_error(format!("a bytes-like object is required, not {}", sep))
-                );
-            }
-        };
+    pub fn partition(&self, sep: &PyByteInner, reverse: bool) -> PyResult<(Vec<u8>, Vec<u8>)> {
         let splitted = if reverse {
-            split_slice_reverse(&self.elements, &sep, 1)
+            split_slice_reverse(&self.elements, &sep.elements, 1)
         } else {
-            split_slice(&self.elements, &sep, 1)
+            split_slice(&self.elements, &sep.elements, 1)
         };
         Ok((splitted[0].to_vec(), splitted[1].to_vec()))
     }
@@ -890,17 +876,13 @@ impl PyByteInner {
         res
     }
 
-    pub fn splitlines(
-        &self,
-        options: ByteInnerSplitlinesOptions,
-        vm: &VirtualMachine,
-    ) -> PyResult<Vec<&[u8]>> {
-        let keepends = options.get_value(vm)?;
+    pub fn splitlines(&self, options: ByteInnerSplitlinesOptions) -> Vec<&[u8]> {
+        let keepends = options.get_value();
 
         let mut res = vec![];
 
         if self.elements.is_empty() {
-            return Ok(vec![]);
+            return vec![];
         }
 
         let mut prev_index = 0;
@@ -935,7 +917,7 @@ impl PyByteInner {
             }
         }
 
-        Ok(res)
+        res
     }
 
     pub fn zfill(&self, width: PyIntRef) -> Vec<u8> {
@@ -960,37 +942,16 @@ impl PyByteInner {
 
     pub fn replace(
         &self,
-        old: PyObjectRef,
-        new: PyObjectRef,
+        old: PyByteInner,
+        new: PyByteInner,
         count: OptionalArg<PyIntRef>,
-        vm: &VirtualMachine,
     ) -> PyResult<Vec<u8>> {
-        let old = match try_as_bytes_like(&old) {
-            Some(value) => value,
-            None => {
-                return Err(
-                    vm.new_type_error(format!("a bytes-like object is required, not {}", old))
-                );
-            }
-        };
-
-        let new = match try_as_bytes_like(&new) {
-            Some(value) => value,
-            None => {
-                return Err(
-                    vm.new_type_error(format!("a bytes-like object is required, not {}", new))
-                );
-            }
-        };
-
-        let count = if let OptionalArg::Present(int) = count {
-            if let Some(value) = int.as_bigint().to_u32() {
-                value
-            } else {
-                self.elements.len() as u32
-            }
-        } else {
-            self.elements.len() as u32
+        let count = match count.into_option() {
+            Some(int) => int
+                .as_bigint()
+                .to_u32()
+                .unwrap_or(self.elements.len() as u32),
+            None => self.elements.len() as u32,
         };
 
         let mut res = vec![];
@@ -1003,8 +964,8 @@ impl PyByteInner {
                 res.extend_from_slice(&slice[index..]);
                 break;
             }
-            if &slice[index..index + old.len()] == old.as_slice() {
-                res.extend_from_slice(&new);
+            if &slice[index..index + old.len()] == old.elements.as_slice() {
+                res.extend_from_slice(&new.elements);
                 index += old.len();
                 done += 1;
             } else {
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -1,4 +1,5 @@
 use crate::obj::objint::PyIntRef;
+
 use crate::obj::objslice::PySliceRef;
 use crate::obj::objstr::PyStringRef;
 use crate::obj::objtuple::PyTupleRef;
@@ -9,7 +10,9 @@ use core::cell::Cell;
 use std::ops::Deref;
 
 use crate::function::OptionalArg;
-use crate::pyobject::{PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue};
+use crate::pyobject::{
+    PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
+};
 
 use super::objbyteinner::{
     ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
@@ -360,17 +363,21 @@ impl PyBytesRef {
     #[pymethod(name = "partition")]
     fn partition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         // TODO:  when implementing bytearray,remember sep ALWAYS converted to  bytearray
-        // even it's bytes or memoryview
-        let (left, right) = self.inner.partition(&sep, false, vm)?;
+        // even it's bytes or memoryview so PyByteInner wiil be ok for args
+        let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
+
+        let (left, right) = self.inner.partition(&sepa, false)?;
         Ok(vm
             .ctx
             .new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
     }
     #[pymethod(name = "rpartition")]
     fn rpartition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
         // TODO:  when implementing bytearray,remember sep ALWAYS converted to  bytearray
-        // even it's bytes or memoryview
-        let (left, right) = self.inner.partition(&sep, true, vm)?;
+        // even it's bytes or memoryview so PyByteInner wiil be ok for args
+        let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
+
+        let (left, right) = self.inner.partition(&sepa, true)?;
         Ok(vm
             .ctx
             .new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
@@ -385,7 +392,7 @@ impl PyBytesRef {
     fn splitlines(self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
         let as_bytes = self
             .inner
-            .splitlines(options, vm)?
+            .splitlines(options)
             .iter()
             .map(|x| vm.ctx.new_bytes(x.to_vec()))
             .collect::<Vec<PyObjectRef>>();
@@ -400,12 +407,12 @@ impl PyBytesRef {
     #[pymethod(name = "replace")]
     fn replace(
         self,
-        old: PyObjectRef,
-        new: PyObjectRef,
+        old: PyByteInner,
+        new: PyByteInner,
         count: OptionalArg<PyIntRef>,
         vm: &VirtualMachine,
     ) -> PyResult {
-        Ok(vm.ctx.new_bytes(self.inner.replace(old, new, count, vm)?))
+        Ok(vm.ctx.new_bytes(self.inner.replace(old, new, count)?))
     }
 }
 
