Move SplitLinesArgs and ExpandTabsArg to pystr

youknowone · youknowone · commit 14f1a4294de3 · 2020-04-19T04:23:32.000+09:00
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
@@ -4,16 +4,15 @@ use std::convert::TryFrom;
 use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
 
 use super::objbyteinner::{
-    ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
-    ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, ByteOr,
-    PyByteInner,
+    ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions,
+    ByteInnerTranslateOptions, ByteOr, PyByteInner,
 };
 use super::objint::PyIntRef;
 use super::objiter;
 use super::objslice::PySliceRef;
 use super::objstr::{PyString, PyStringRef};
 use super::objtype::PyClassRef;
-use super::pystr::PyCommonString;
+use super::pystr::{self, PyCommonString};
 use crate::cformat::CFormatString;
 use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objstr::do_cformat_string;
@@ -440,12 +439,12 @@ impl PyByteArray {
     }
 
     #[pymethod(name = "expandtabs")]
-    fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> PyByteArray {
+    fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> PyByteArray {
         self.borrow_value().expandtabs(options).into()
     }
 
     #[pymethod(name = "splitlines")]
-    fn splitlines(&self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
+    fn splitlines(&self, options: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyResult {
         let as_bytes = self
             .borrow_value()
             .splitlines(options)
diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs
@@ -257,21 +257,6 @@ impl ByteInnerTranslateOptions {
 
 pub type ByteInnerSplitOptions = pystr::SplitArgs<PyByteInner, [u8], u8>;
 
-#[derive(FromArgs)]
-pub struct ByteInnerExpandtabsOptions {
-    #[pyarg(positional_or_keyword, optional = true)]
-    tabsize: OptionalArg<PyIntRef>,
-}
-
-impl ByteInnerExpandtabsOptions {
-    pub fn get_value(self) -> usize {
-        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)]
@@ -1014,8 +999,8 @@ impl PyByteInner {
         Ok((front, has_mid, back))
     }
 
-    pub fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> Vec<u8> {
-        let tabsize = options.get_value();
+    pub fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> Vec<u8> {
+        let tabsize = options.tabsize();
         let mut counter: usize = 0;
         let mut res = vec![];
 
@@ -1046,9 +1031,7 @@ impl PyByteInner {
         res
     }
 
-    pub fn splitlines(&self, options: ByteInnerSplitlinesOptions) -> Vec<&[u8]> {
-        let keepends = options.get_value();
-
+    pub fn splitlines(&self, options: pystr::SplitLinesArgs) -> Vec<&[u8]> {
         let mut res = vec![];
 
         if self.elements.is_empty() {
@@ -1057,7 +1040,7 @@ impl PyByteInner {
 
         let mut prev_index = 0;
         let mut index = 0;
-        let keep = if keepends { 1 } else { 0 };
+        let keep = if options.keepends { 1 } else { 0 };
         let slice = &self.elements;
 
         while index < slice.len() {
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -4,15 +4,15 @@ use std::ops::Deref;
 use std::str::FromStr;
 
 use super::objbyteinner::{
-    ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
-    ByteInnerSplitOptions, ByteInnerSplitlinesOptions, ByteInnerTranslateOptions, PyByteInner,
+    ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerSplitOptions,
+    ByteInnerTranslateOptions, PyByteInner,
 };
 use super::objint::PyIntRef;
 use super::objiter;
 use super::objslice::PySliceRef;
 use super::objstr::{PyString, PyStringRef};
 use super::objtype::PyClassRef;
-use super::pystr::PyCommonString;
+use super::pystr::{self, PyCommonString};
 use crate::cformat::CFormatString;
 use crate::function::{OptionalArg, OptionalOption};
 use crate::obj::objstr::do_cformat_string;
@@ -401,12 +401,12 @@ impl PyBytes {
     }
 
     #[pymethod(name = "expandtabs")]
-    fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> PyBytes {
+    fn expandtabs(&self, options: pystr::ExpandTabsArgs) -> PyBytes {
         self.inner.expandtabs(options).into()
     }
 
     #[pymethod(name = "splitlines")]
-    fn splitlines(&self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
+    fn splitlines(&self, options: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyResult {
         let as_bytes = self
             .inner
             .splitlines(options)
diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs
@@ -168,18 +168,6 @@ struct StrArgs {
     errors: OptionalArg<PyStringRef>,
 }
 
-#[derive(FromArgs)]
-struct SplitLineArgs {
-    #[pyarg(positional_or_keyword, optional = true)]
-    keepends: OptionalArg<bool>,
-}
-
-#[derive(FromArgs)]
-struct ExpandtabsArgs {
-    #[pyarg(positional_or_keyword, optional = true)]
-    tabsize: OptionalArg<usize>,
-}
-
 #[pyimpl(flags(BASETYPE))]
 impl PyString {
     #[pyslot]
@@ -797,14 +785,13 @@ impl PyString {
     }
 
     #[pymethod]
-    fn splitlines(&self, args: SplitLineArgs, vm: &VirtualMachine) -> PyObjectRef {
-        let keepends = args.keepends.unwrap_or(false);
+    fn splitlines(&self, args: pystr::SplitLinesArgs, vm: &VirtualMachine) -> PyObjectRef {
         let mut elements = vec![];
         let mut curr = "".to_owned();
         let mut chars = self.value.chars().peekable();
         while let Some(ch) = chars.next() {
             if ch == '\n' || ch == '\r' {
-                if keepends {
+                if args.keepends {
                     curr.push(ch);
                 }
                 if ch == '\r' && chars.peek() == Some(&'\n') {
@@ -1080,8 +1067,8 @@ impl PyString {
     }
 
     #[pymethod]
-    fn expandtabs(&self, args: ExpandtabsArgs) -> String {
-        let tab_stop = args.tabsize.unwrap_or(8);
+    fn expandtabs(&self, args: pystr::ExpandTabsArgs) -> String {
+        let tab_stop = args.tabsize();
         let mut expanded_str = String::with_capacity(self.value.len());
         let mut tab_size = tab_stop;
         let mut col_count = 0 as usize;
diff --git a/vm/src/obj/pystr.rs b/vm/src/obj/pystr.rs
@@ -1,6 +1,7 @@
 use crate::function::{single_or_tuple_any, OptionalOption};
 use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
 use crate::vm::VirtualMachine;
+use num_traits::cast::ToPrimitive;
 
 #[derive(FromArgs)]
 pub struct SplitArgs<T, S, E>
@@ -35,6 +36,24 @@ where
     }
 }
 
+#[derive(FromArgs)]
+pub struct SplitLinesArgs {
+    #[pyarg(positional_or_keyword, default = "false")]
+    pub keepends: bool,
+}
+
+#[derive(FromArgs)]
+pub struct ExpandTabsArgs {
+    #[pyarg(positional_or_keyword, default = "8")]
+    tabsize: isize,
+}
+
+impl ExpandTabsArgs {
+    pub fn tabsize(&self) -> usize {
+        self.tabsize.to_usize().unwrap_or(0)
+    }
+}
+
 // help get optional string indices
 pub fn adjust_indices(
     start: OptionalOption<isize>,
