Skip to content

Commit 33c8551

Browse files
youknowoneSnowapril
youknowone
authored and
Snowapril
committed
Optional subrange operation in StartsEndsWithArgs
1 parent dc5debe commit 33c8551

File tree

4 files changed

+68
-68
lines changed

4 files changed

+68
-68
lines changed

‎vm/src/anystr.rs

Lines changed: 23 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -8,19 +8,6 @@ use crate::{
88
use num_traits::{cast::ToPrimitive, sign::Signed};
99
use std::str::FromStr;
1010

11-
pub enum AnyStrRange {
12-
CharsRange(std::ops::Range<usize>),
13-
BytesRange(std::ops::Range<usize>),
14-
}
15-
16-
impl AnyStrRange {
17-
pub fn is_normal(&self) -> bool {
18-
match self {
19-
AnyStrRange::CharsRange(range) | AnyStrRange::BytesRange(range) => range.is_normal(),
20-
}
21-
}
22-
}
23-
2411
#[derive(FromArgs)]
2512
pub struct SplitArgs<'s, T: TryFromObject + AnyStrWrapper<'s>> {
2613
#[pyarg(any, default)]
@@ -74,13 +61,31 @@ pub struct StartsEndsWithArgs {
7461
}
7562

7663
impl StartsEndsWithArgs {
77-
pub fn get_value(self, len: usize) -> (PyObjectRef, std::ops::Range<usize>) {
78-
let range = adjust_indices(self.start, self.end, len);
64+
pub fn get_value(self, len: usize) -> (PyObjectRef, Option<std::ops::Range<usize>>) {
65+
let range = if self.start.is_some() || self.end.is_some() {
66+
Some(adjust_indices(self.start, self.end, len))
67+
} else {
68+
None
69+
};
7970
(self.affix, range)
8071
}
8172

82-
pub fn has_subrange(&self) -> bool {
83-
self.start.is_some() || self.end.is_some()
73+
#[inline]
74+
pub fn prepare<'s, S, F>(self, s: &'s S, len: usize, substr: F) -> Option<(PyObjectRef, &'s S)>
75+
where
76+
S: ?Sized + AnyStr<'s>,
77+
F: Fn(&'s S, std::ops::Range<usize>) -> &'s S,
78+
{
79+
let (affix, range) = self.get_value(len);
80+
let substr = if let Some(range) = range {
81+
if !range.is_normal() {
82+
return None;
83+
}
84+
substr(s, range)
85+
} else {
86+
s
87+
};
88+
Some((affix, substr))
8489
}
8590
}
8691

@@ -211,7 +216,6 @@ pub trait AnyStr<'s>: 's {
211216
fn py_startsendswith<T, F>(
212217
&self,
213218
affix: PyObjectRef,
214-
range: AnyStrRange,
215219
func_name: &str,
216220
py_type_name: &str,
217221
func: F,
@@ -221,17 +225,9 @@ pub trait AnyStr<'s>: 's {
221225
T: TryFromObject,
222226
F: Fn(&Self, &T) -> bool,
223227
{
224-
if !range.is_normal() {
225-
return Ok(false);
226-
}
227-
let value = match range {
228-
AnyStrRange::BytesRange(range) => self.get_bytes(range),
229-
AnyStrRange::CharsRange(range) => self.get_chars(range),
230-
};
231-
232228
single_or_tuple_any(
233229
affix,
234-
&|s: &T| Ok(func(value, s)),
230+
&|s: &T| Ok(func(self, s)),
235231
&|o| {
236232
format!(
237233
"{} first arg must be {} or a tuple of {}, not {}",

‎vm/src/builtins/bytearray.rs

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -435,10 +435,14 @@ impl PyByteArray {
435435

436436
#[pymethod]
437437
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
438-
let (affix, range) = options.get_value(self.len());
439-
self.borrow_buf().py_startsendswith(
438+
let borrowed = self.borrow_buf();
439+
let (affix, substr) =
440+
match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
441+
Some(x) => x,
442+
None => return Ok(false),
443+
};
444+
substr.py_startsendswith(
440445
affix,
441-
anystr::AnyStrRange::BytesRange(range),
442446
"endswith",
443447
"bytes",
444448
|s, x: &PyBytesInner| s.ends_with(&x.elements[..]),
@@ -452,10 +456,14 @@ impl PyByteArray {
452456
options: anystr::StartsEndsWithArgs,
453457
vm: &VirtualMachine,
454458
) -> PyResult<bool> {
455-
let (affix, range) = options.get_value(self.len());
456-
self.borrow_buf().py_startsendswith(
459+
let borrowed = self.borrow_buf();
460+
let (affix, substr) =
461+
match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
462+
Some(x) => x,
463+
None => return Ok(false),
464+
};
465+
substr.py_startsendswith(
457466
affix,
458-
anystr::AnyStrRange::BytesRange(range),
459467
"startswith",
460468
"bytes",
461469
|s, x: &PyBytesInner| s.starts_with(&x.elements[..]),

‎vm/src/builtins/bytes.rs

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -275,10 +275,13 @@ impl PyBytes {
275275

276276
#[pymethod]
277277
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
278-
let (affix, range) = options.get_value(self.len());
279-
self.inner.elements[..].py_startsendswith(
278+
let (affix, substr) =
279+
match options.prepare(&self.inner.elements[..], self.len(), |s, r| s.get_bytes(r)) {
280+
Some(x) => x,
281+
None => return Ok(false),
282+
};
283+
substr.py_startsendswith(
280284
affix,
281-
anystr::AnyStrRange::BytesRange(range),
282285
"endswith",
283286
"bytes",
284287
|s, x: &PyBytesInner| s.ends_with(&x.elements[..]),
@@ -292,10 +295,13 @@ impl PyBytes {
292295
options: anystr::StartsEndsWithArgs,
293296
vm: &VirtualMachine,
294297
) -> PyResult<bool> {
295-
let (affix, range) = options.get_value(self.len());
296-
self.inner.elements[..].py_startsendswith(
298+
let (affix, substr) =
299+
match options.prepare(&self.inner.elements[..], self.len(), |s, r| s.get_bytes(r)) {
300+
Some(x) => x,
301+
None => return Ok(false),
302+
};
303+
substr.py_startsendswith(
297304
affix,
298-
anystr::AnyStrRange::BytesRange(range),
299305
"startswith",
300306
"bytes",
301307
|s, x: &PyBytesInner| s.starts_with(&x.elements[..]),

‎vm/src/builtins/pystr.rs

Lines changed: 19 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use super::{
44
PositionIterInternal, PyBytesRef, PyDict, PyTupleRef, PyTypeRef,
55
};
66
use crate::{
7-
anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrRange, AnyStrWrapper},
7+
anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrWrapper},
88
format::{FormatSpec, FormatString, FromTemplate},
99
function::{ArgIterable, FuncArgs, IntoPyException, IntoPyObject, OptionalArg, OptionalOption},
1010
protocol::PyIterReturn,
@@ -705,21 +705,14 @@ impl PyStr {
705705
}
706706

707707
#[pymethod]
708-
fn endswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
709-
let has_subrange = args.has_subrange();
710-
let len = if args.has_subrange() {
711-
self.char_len()
712-
} else {
713-
self.byte_len()
714-
};
715-
let (affix, range) = args.get_value(len);
716-
self.as_str().py_startsendswith(
708+
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
709+
let (affix, substr) =
710+
match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
711+
Some(x) => x,
712+
None => return Ok(false),
713+
};
714+
substr.py_startsendswith(
717715
affix,
718-
if has_subrange {
719-
AnyStrRange::CharsRange(range)
720-
} else {
721-
AnyStrRange::BytesRange(range)
722-
},
723716
"endswith",
724717
"str",
725718
|s, x: &PyStrRef| s.ends_with(x.as_str()),
@@ -728,21 +721,18 @@ impl PyStr {
728721
}
729722

730723
#[pymethod]
731-
fn startswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
732-
let has_subrange = args.has_subrange();
733-
let len = if has_subrange {
734-
self.char_len()
735-
} else {
736-
self.byte_len()
737-
};
738-
let (affix, range) = args.get_value(len);
739-
self.as_str().py_startsendswith(
724+
fn startswith(
725+
&self,
726+
options: anystr::StartsEndsWithArgs,
727+
vm: &VirtualMachine,
728+
) -> PyResult<bool> {
729+
let (affix, substr) =
730+
match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
731+
Some(x) => x,
732+
None => return Ok(false),
733+
};
734+
substr.py_startsendswith(
740735
affix,
741-
if has_subrange {
742-
AnyStrRange::CharsRange(range)
743-
} else {
744-
AnyStrRange::BytesRange(range)
745-
},
746736
"startswith",
747737
"str",
748738
|s, x: &PyStrRef| s.starts_with(x.as_str()),

0 commit comments

Comments
 (0)