Skip to content

Commit b7a399e

Browse files
youknowoneyouknowone
committed
merge SplitArgs from PyByte* and PyString
1 parent 162e49f commit b7a399e

File tree

3 files changed

+83
-76
lines changed

3 files changed

+83
-76
lines changed

‎vm/src/obj/objbyteinner.rs

Lines changed: 16 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use super::objnone::PyNoneRef;
1414
use super::objsequence::PySliceableSequence;
1515
use super::objslice::PySliceRef;
1616
use super::objstr::{self, PyString, PyStringRef};
17-
use super::pystr::{self, PyCommonString, StringRange};
17+
use super::pystr::{self, PyCommonString, PyCommonStringWrapper, StringRange};
1818
use crate::function::{OptionalArg, OptionalOption};
1919
use crate::pyhash;
2020
use crate::pyobject::{
@@ -255,28 +255,7 @@ impl ByteInnerTranslateOptions {
255255
}
256256
}
257257

258-
#[derive(FromArgs)]
259-
pub struct ByteInnerSplitOptions {
260-
#[pyarg(positional_or_keyword, default = "None")]
261-
sep: Option<PyByteInner>,
262-
#[pyarg(positional_or_keyword, default = "-1")]
263-
maxsplit: isize,
264-
}
265-
266-
impl ByteInnerSplitOptions {
267-
pub fn get_value(self, vm: &VirtualMachine) -> PyResult<(Option<Vec<u8>>, isize)> {
268-
let sep = if let Some(s) = self.sep {
269-
let sep = s.elements;
270-
if sep.is_empty() {
271-
return Err(vm.new_value_error("empty separator".to_owned()));
272-
}
273-
Some(sep)
274-
} else {
275-
None
276-
};
277-
Ok((sep, self.maxsplit))
278-
}
279-
}
258+
pub type ByteInnerSplitOptions = pystr::SplitArgs<PyByteInner, [u8], u8>;
280259

281260
#[derive(FromArgs)]
282261
pub struct ByteInnerExpandtabsOptions {
@@ -967,19 +946,13 @@ impl PyByteInner {
967946
where
968947
F: Fn(&[u8], &VirtualMachine) -> PyObjectRef,
969948
{
970-
let (sep, maxsplit) = options.get_value(vm)?;
971-
let sep_ref = match sep {
972-
Some(ref v) => Some(&v[..]),
973-
None => None,
974-
};
975949
let elements = self.elements.py_split(
976-
sep_ref,
977-
maxsplit,
950+
options,
978951
vm,
979952
|v, s, vm| v.split_str(s).map(|v| convert(v, vm)).collect(),
980953
|v, s, n, vm| v.splitn_str(n, s).map(|v| convert(v, vm)).collect(),
981954
|v, n, vm| v.py_split_whitespace(n, |v| convert(v, vm)),
982-
);
955+
)?;
983956
Ok(vm.ctx.new_list(elements))
984957
}
985958

@@ -992,19 +965,13 @@ impl PyByteInner {
992965
where
993966
F: Fn(&[u8], &VirtualMachine) -> PyObjectRef,
994967
{
995-
let (sep, maxsplit) = options.get_value(vm)?;
996-
let sep_ref = match sep {
997-
Some(ref v) => Some(&v[..]),
998-
None => None,
999-
};
1000968
let mut elements = self.elements.py_split(
1001-
sep_ref,
1002-
maxsplit,
969+
options,
1003970
vm,
1004971
|v, s, vm| v.rsplit_str(s).map(|v| convert(v, vm)).collect(),
1005972
|v, s, n, vm| v.rsplitn_str(n, s).map(|v| convert(v, vm)).collect(),
1006973
|v, n, vm| v.py_rsplit_whitespace(n, |v| convert(v, vm)),
1007-
);
974+
)?;
1008975
elements.reverse();
1009976
Ok(vm.ctx.new_list(elements))
1010977
}
@@ -1297,13 +1264,23 @@ pub fn bytes_zfill(bytes: &[u8], width: usize) -> Vec<u8> {
12971264
}
12981265
}
12991266

1267+
impl PyCommonStringWrapper<[u8]> for PyByteInner {
1268+
fn as_ref(&self) -> &[u8] {
1269+
&self.elements
1270+
}
1271+
}
1272+
13001273
const ASCII_WHITESPACES: [u8; 6] = [0x20, 0x09, 0x0a, 0x0c, 0x0d, 0x0b];
13011274

13021275
impl PyCommonString<u8> for [u8] {
13031276
fn get_slice(&self, range: std::ops::Range<usize>) -> &Self {
13041277
&self[range]
13051278
}
13061279

1280+
fn is_empty(&self) -> bool {
1281+
Self::is_empty(self)
1282+
}
1283+
13071284
fn len(&self) -> usize {
13081285
Self::len(self)
13091286
}

‎vm/src/obj/objstr.rs

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ use super::objsequence::PySliceableSequence;
2121
use super::objslice::PySliceRef;
2222
use super::objtuple;
2323
use super::objtype::{self, PyClassRef};
24-
use super::pystr::{adjust_indices, PyCommonString, StringRange};
24+
use super::pystr::{self, adjust_indices, PyCommonString, PyCommonStringWrapper, StringRange};
2525
use crate::cformat::{
2626
CFormatPart, CFormatPreconversor, CFormatQuantity, CFormatSpec, CFormatString, CFormatType,
2727
CNumberType,
@@ -457,26 +457,24 @@ impl PyString {
457457
#[pymethod]
458458
fn split(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult {
459459
let elements = self.value.py_split(
460-
args.non_empty_sep(vm)?,
461-
args.maxsplit,
460+
args,
462461
vm,
463462
|v, s, vm| v.split(s).map(|s| vm.ctx.new_str(s)).collect(),
464463
|v, s, n, vm| v.splitn(n, s).map(|s| vm.ctx.new_str(s)).collect(),
465464
|v, n, vm| v.py_split_whitespace(n, |s| vm.ctx.new_str(s)),
466-
);
465+
)?;
467466
Ok(vm.ctx.new_list(elements))
468467
}
469468

470469
#[pymethod]
471470
fn rsplit(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult {
472471
let mut elements = self.value.py_split(
473-
args.non_empty_sep(vm)?,
474-
args.maxsplit,
472+
args,
475473
vm,
476474
|v, s, vm| v.rsplit(s).map(|s| vm.ctx.new_str(s)).collect(),
477475
|v, s, n, vm| v.rsplitn(n, s).map(|s| vm.ctx.new_str(s)).collect(),
478476
|v, n, vm| v.py_rsplit_whitespace(n, |s| vm.ctx.new_str(s)),
479-
);
477+
)?;
480478
// Unlike Python rsplit, Rust rsplitn returns an iterator that
481479
// starts from the end of the string.
482480
elements.reverse();
@@ -1298,28 +1296,7 @@ impl TryFromObject for std::ffi::CString {
12981296
}
12991297
}
13001298

1301-
#[derive(FromArgs)]
1302-
struct SplitArgs {
1303-
#[pyarg(positional_or_keyword, default = "None")]
1304-
sep: Option<PyStringRef>,
1305-
#[pyarg(positional_or_keyword, default = "-1")]
1306-
maxsplit: isize,
1307-
}
1308-
1309-
impl SplitArgs {
1310-
fn non_empty_sep<'a>(&'a self, vm: &VirtualMachine) -> PyResult<Option<&'a str>> {
1311-
let sep = if let Some(s) = self.sep.as_ref() {
1312-
let sep = s.as_str();
1313-
if sep.is_empty() {
1314-
return Err(vm.new_value_error("empty separator".to_owned()));
1315-
}
1316-
Some(sep)
1317-
} else {
1318-
None
1319-
};
1320-
Ok(sep)
1321-
}
1322-
}
1299+
type SplitArgs = pystr::SplitArgs<PyStringRef, str, char>;
13231300

13241301
pub fn init(ctx: &PyContext) {
13251302
PyString::extend_class(ctx, &ctx.types.str_type);
@@ -1799,11 +1776,21 @@ mod tests {
17991776
}
18001777
}
18011778

1779+
impl PyCommonStringWrapper<str> for PyStringRef {
1780+
fn as_ref(&self) -> &str {
1781+
self.value.as_str()
1782+
}
1783+
}
1784+
18021785
impl PyCommonString<char> for str {
18031786
fn get_slice(&self, range: std::ops::Range<usize>) -> &Self {
18041787
&self[range]
18051788
}
18061789

1790+
fn is_empty(&self) -> bool {
1791+
Self::is_empty(self)
1792+
}
1793+
18071794
fn len(&self) -> usize {
18081795
Self::len(self)
18091796
}

‎vm/src/obj/pystr.rs

Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,39 @@ use crate::function::{single_or_tuple_any, OptionalOption};
22
use crate::pyobject::{PyObjectRef, PyResult, TryFromObject, TypeProtocol};
33
use crate::vm::VirtualMachine;
44

5+
#[derive(FromArgs)]
6+
pub struct SplitArgs<T, S, E>
7+
where
8+
T: TryFromObject + PyCommonStringWrapper<S>,
9+
S: ?Sized + PyCommonString<E>,
10+
{
11+
#[pyarg(positional_or_keyword, default = "None")]
12+
sep: Option<T>,
13+
#[pyarg(positional_or_keyword, default = "-1")]
14+
maxsplit: isize,
15+
_phantom1: std::marker::PhantomData<S>,
16+
_phantom2: std::marker::PhantomData<E>,
17+
}
18+
19+
impl<T, S, E> SplitArgs<T, S, E>
20+
where
21+
T: TryFromObject + PyCommonStringWrapper<S>,
22+
S: ?Sized + PyCommonString<E>,
23+
{
24+
pub fn get_value(self, vm: &VirtualMachine) -> PyResult<(Option<T>, isize)> {
25+
let sep = if let Some(s) = self.sep {
26+
let sep = s.as_ref();
27+
if sep.is_empty() {
28+
return Err(vm.new_value_error("empty separator".to_owned()));
29+
}
30+
Some(s)
31+
} else {
32+
None
33+
};
34+
Ok((sep, self.maxsplit))
35+
}
36+
}
37+
538
// help get optional string indices
639
pub fn adjust_indices(
740
start: OptionalOption<isize>,
@@ -37,33 +70,43 @@ impl StringRange for std::ops::Range<usize> {
3770
}
3871
}
3972

73+
pub trait PyCommonStringWrapper<S>
74+
where
75+
S: ?Sized,
76+
{
77+
fn as_ref(&self) -> &S;
78+
}
79+
4080
pub trait PyCommonString<E> {
4181
fn get_slice(&self, range: std::ops::Range<usize>) -> &Self;
4282
fn len(&self) -> usize;
83+
fn is_empty(&self) -> bool;
4384

44-
fn py_split<SP, SN, SW, R>(
85+
fn py_split<T, SP, SN, SW, R>(
4586
&self,
46-
sep: Option<&Self>,
47-
maxsplit: isize,
87+
args: SplitArgs<T, Self, E>,
4888
vm: &VirtualMachine,
4989
split: SP,
5090
splitn: SN,
5191
splitw: SW,
52-
) -> Vec<R>
92+
) -> PyResult<Vec<R>>
5393
where
94+
T: TryFromObject + PyCommonStringWrapper<Self>,
5495
SP: Fn(&Self, &Self, &VirtualMachine) -> Vec<R>,
5596
SN: Fn(&Self, &Self, usize, &VirtualMachine) -> Vec<R>,
5697
SW: Fn(&Self, isize, &VirtualMachine) -> Vec<R>,
5798
{
58-
if let Some(pattern) = sep {
99+
let (sep, maxsplit) = args.get_value(vm)?;
100+
let splited = if let Some(pattern) = sep {
59101
if maxsplit < 0 {
60-
split(self, pattern, vm)
102+
split(self, pattern.as_ref(), vm)
61103
} else {
62-
splitn(self, pattern, (maxsplit + 1) as usize, vm)
104+
splitn(self, pattern.as_ref(), (maxsplit + 1) as usize, vm)
63105
}
64106
} else {
65107
splitw(self, maxsplit, vm)
66-
}
108+
};
109+
Ok(splited)
67110
}
68111
fn py_split_whitespace<F>(&self, maxsplit: isize, convert: F) -> Vec<PyObjectRef>
69112
where

0 commit comments

Comments
 (0)