Skip to content

Commit 319914d

Browse files
youknowoneyouknowone
committed
Optional subrange operation in StartsEndsWithArgs
1 parent a18e6a7 commit 319914d

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
@@ -430,10 +430,14 @@ impl PyByteArray {
430430

431431
#[pymethod]
432432
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
433-
let (affix, range) = options.get_value(self.len());
434-
self.borrow_buf().py_startsendswith(
433+
let borrowed = self.borrow_buf();
434+
let (affix, substr) =
435+
match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
436+
Some(x) => x,
437+
None => return Ok(false),
438+
};
439+
substr.py_startsendswith(
435440
affix,
436-
anystr::AnyStrRange::BytesRange(range),
437441
"endswith",
438442
"bytes",
439443
|s, x: &PyBytesInner| s.ends_with(&x.elements[..]),
@@ -447,10 +451,14 @@ impl PyByteArray {
447451
options: anystr::StartsEndsWithArgs,
448452
vm: &VirtualMachine,
449453
) -> PyResult<bool> {
450-
let (affix, range) = options.get_value(self.len());
451-
self.borrow_buf().py_startsendswith(
454+
let borrowed = self.borrow_buf();
455+
let (affix, substr) =
456+
match options.prepare(&*borrowed, borrowed.len(), |s, r| s.get_bytes(r)) {
457+
Some(x) => x,
458+
None => return Ok(false),
459+
};
460+
substr.py_startsendswith(
452461
affix,
453-
anystr::AnyStrRange::BytesRange(range),
454462
"startswith",
455463
"bytes",
456464
|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,
@@ -687,21 +687,14 @@ impl PyStr {
687687
}
688688

689689
#[pymethod]
690-
fn endswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
691-
let has_subrange = args.has_subrange();
692-
let len = if args.has_subrange() {
693-
self.char_len()
694-
} else {
695-
self.byte_len()
696-
};
697-
let (affix, range) = args.get_value(len);
698-
self.as_str().py_startsendswith(
690+
fn endswith(&self, options: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
691+
let (affix, substr) =
692+
match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
693+
Some(x) => x,
694+
None => return Ok(false),
695+
};
696+
substr.py_startsendswith(
699697
affix,
700-
if has_subrange {
701-
AnyStrRange::CharsRange(range)
702-
} else {
703-
AnyStrRange::BytesRange(range)
704-
},
705698
"endswith",
706699
"str",
707700
|s, x: &PyStrRef| s.ends_with(x.as_str()),
@@ -710,21 +703,18 @@ impl PyStr {
710703
}
711704

712705
#[pymethod]
713-
fn startswith(&self, args: anystr::StartsEndsWithArgs, vm: &VirtualMachine) -> PyResult<bool> {
714-
let has_subrange = args.has_subrange();
715-
let len = if has_subrange {
716-
self.char_len()
717-
} else {
718-
self.byte_len()
719-
};
720-
let (affix, range) = args.get_value(len);
721-
self.as_str().py_startsendswith(
706+
fn startswith(
707+
&self,
708+
options: anystr::StartsEndsWithArgs,
709+
vm: &VirtualMachine,
710+
) -> PyResult<bool> {
711+
let (affix, substr) =
712+
match options.prepare(self.as_str(), self.len(), |s, r| s.get_chars(r)) {
713+
Some(x) => x,
714+
None => return Ok(false),
715+
};
716+
substr.py_startsendswith(
722717
affix,
723-
if has_subrange {
724-
AnyStrRange::CharsRange(range)
725-
} else {
726-
AnyStrRange::BytesRange(range)
727-
},
728718
"startswith",
729719
"str",
730720
|s, x: &PyStrRef| s.starts_with(x.as_str()),

0 commit comments

Comments
 (0)