Skip to content

Commit 4c9486f

Browse files
youknowoneyouknowone
committed
Fix bytes.{start|end}swith
1 parent 0d2b817 commit 4c9486f

File tree

6 files changed

+169
-143
lines changed

6 files changed

+169
-143
lines changed

‎Lib/test/test_bytes.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -548,8 +548,6 @@ def test_count(self):
548548
self.assertEqual(b.count(i, 1, 3), 1)
549549
self.assertEqual(b.count(p, 7, 9), 1)
550550

551-
# TODO: RUSTPYTHON
552-
@unittest.expectedFailure
553551
def test_startswith(self):
554552
b = self.type2test(b'hello')
555553
self.assertFalse(self.type2test().startswith(b"anything"))
@@ -564,8 +562,6 @@ def test_startswith(self):
564562
self.assertIn('bytes', exc)
565563
self.assertIn('tuple', exc)
566564

567-
# TODO: RUSTPYTHON
568-
@unittest.expectedFailure
569565
def test_endswith(self):
570566
b = self.type2test(b'hello')
571567
self.assertFalse(bytearray().endswith(b"anything"))

‎vm/src/obj/objbytearray.rs

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ use super::objint::PyIntRef;
1212
use super::objiter;
1313
use super::objslice::PySliceRef;
1414
use super::objstr::{PyString, PyStringRef};
15-
use super::objtuple::PyTupleRef;
1615
use super::objtype::PyClassRef;
16+
use super::pystr::PyCommonString;
1717
use crate::cformat::CFormatString;
1818
use crate::function::{OptionalArg, OptionalOption};
1919
use crate::obj::objstr::do_cformat_string;
@@ -303,25 +303,39 @@ impl PyByteArray {
303303
#[pymethod(name = "endswith")]
304304
fn endswith(
305305
&self,
306-
suffix: Either<PyByteInner, PyTupleRef>,
307-
start: OptionalArg<PyObjectRef>,
308-
end: OptionalArg<PyObjectRef>,
306+
suffix: PyObjectRef,
307+
start: OptionalArg<Option<isize>>,
308+
end: OptionalArg<Option<isize>>,
309309
vm: &VirtualMachine,
310310
) -> PyResult<bool> {
311-
self.borrow_value()
312-
.startsendswith(suffix, start, end, true, vm)
311+
self.borrow_value().elements[..].py_startsendswith(
312+
suffix,
313+
start,
314+
end,
315+
"endswith",
316+
"bytes",
317+
|s, x: &PyByteInner| s.ends_with(&x.elements[..]),
318+
vm,
319+
)
313320
}
314321

315322
#[pymethod(name = "startswith")]
316323
fn startswith(
317324
&self,
318-
prefix: Either<PyByteInner, PyTupleRef>,
319-
start: OptionalArg<PyObjectRef>,
320-
end: OptionalArg<PyObjectRef>,
325+
prefix: PyObjectRef,
326+
start: OptionalArg<Option<isize>>,
327+
end: OptionalArg<Option<isize>>,
321328
vm: &VirtualMachine,
322329
) -> PyResult<bool> {
323-
self.borrow_value()
324-
.startsendswith(prefix, start, end, false, vm)
330+
self.borrow_value().elements[..].py_startsendswith(
331+
prefix,
332+
start,
333+
end,
334+
"startswith",
335+
"bytes",
336+
|s, x: &PyByteInner| s.starts_with(&x.elements[..]),
337+
vm,
338+
)
325339
}
326340

327341
#[pymethod(name = "find")]

‎vm/src/obj/objbyteinner.rs

Lines changed: 12 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,10 @@ use super::objint::{self, PyInt, PyIntRef};
1111
use super::objlist::PyList;
1212
use super::objmemory::PyMemoryView;
1313
use super::objnone::PyNoneRef;
14-
use super::objsequence::{is_valid_slice_arg, PySliceableSequence};
14+
use super::objsequence::PySliceableSequence;
1515
use super::objslice::PySliceRef;
16-
use super::objstr::{self, adjust_indices, PyString, PyStringRef, StringRange};
17-
use super::objtuple::PyTupleRef;
18-
use super::pystr::PyCommonString;
16+
use super::objstr::{self, PyString, PyStringRef};
17+
use super::pystr::{self, PyCommonString, StringRange};
1918
use crate::function::{OptionalArg, OptionalOption};
2019
use crate::pyhash;
2120
use crate::pyobject::{
@@ -172,7 +171,7 @@ impl ByteInnerFindOptions {
172171
Either::A(v) => v.elements.to_vec(),
173172
Either::B(int) => vec![int.as_bigint().byte_or(vm)?],
174173
};
175-
let range = adjust_indices(self.start, self.end, len);
174+
let range = pystr::adjust_indices(self.start, self.end, len);
176175
Ok((sub, range))
177176
}
178177
}
@@ -856,47 +855,6 @@ impl PyByteInner {
856855
Ok(refs)
857856
}
858857

859-
#[inline]
860-
pub fn startsendswith(
861-
&self,
862-
arg: Either<PyByteInner, PyTupleRef>,
863-
start: OptionalArg<PyObjectRef>,
864-
end: OptionalArg<PyObjectRef>,
865-
endswith: bool, // true for endswith, false for startswith
866-
vm: &VirtualMachine,
867-
) -> PyResult<bool> {
868-
let suff = match arg {
869-
Either::A(byte) => byte.elements,
870-
Either::B(tuple) => {
871-
let mut flatten = vec![];
872-
for v in tuple.as_slice() {
873-
flatten.extend(PyByteInner::try_from_object(vm, v.clone())?.elements)
874-
}
875-
flatten
876-
}
877-
};
878-
879-
if suff.is_empty() {
880-
return Ok(true);
881-
}
882-
let range = self.elements.get_slice_range(
883-
&is_valid_slice_arg(start, vm)?,
884-
&is_valid_slice_arg(end, vm)?,
885-
);
886-
887-
if range.end - range.start < suff.len() {
888-
return Ok(false);
889-
}
890-
891-
let offset = if endswith {
892-
(range.end - suff.len())..range.end
893-
} else {
894-
0..suff.len()
895-
};
896-
897-
Ok(suff.as_slice() == &self.elements.do_slice(range)[offset])
898-
}
899-
900858
#[inline]
901859
pub fn find(
902860
&self,
@@ -1342,6 +1300,14 @@ pub fn bytes_zfill(bytes: &[u8], width: usize) -> Vec<u8> {
13421300
const ASCII_WHITESPACES: [u8; 6] = [0x20, 0x09, 0x0a, 0x0c, 0x0d, 0x0b];
13431301

13441302
impl PyCommonString<'_, u8> for [u8] {
1303+
fn get_slice(&self, range: std::ops::Range<usize>) -> &Self {
1304+
&self[range]
1305+
}
1306+
1307+
fn len(&self) -> usize {
1308+
Self::len(self)
1309+
}
1310+
13451311
fn py_split_whitespace<F>(&self, maxsplit: isize, convert: F) -> Vec<PyObjectRef>
13461312
where
13471313
F: Fn(&Self) -> PyObjectRef,

‎vm/src/obj/objbytes.rs

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
use crossbeam_utils::atomic::AtomicCell;
22
use std::mem::size_of;
33
use std::ops::Deref;
4+
use std::str::FromStr;
45

56
use super::objbyteinner::{
67
ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
@@ -10,8 +11,8 @@ use super::objint::PyIntRef;
1011
use super::objiter;
1112
use super::objslice::PySliceRef;
1213
use super::objstr::{PyString, PyStringRef};
13-
use super::objtuple::PyTupleRef;
1414
use super::objtype::PyClassRef;
15+
use super::pystr::PyCommonString;
1516
use crate::cformat::CFormatString;
1617
use crate::function::{OptionalArg, OptionalOption};
1718
use crate::obj::objstr::do_cformat_string;
@@ -23,7 +24,6 @@ use crate::pyobject::{
2324
ThreadSafe, TryFromObject, TypeProtocol,
2425
};
2526
use crate::vm::VirtualMachine;
26-
use std::str::FromStr;
2727

2828
/// "bytes(iterable_of_ints) -> bytes\n\
2929
/// bytes(string, encoding[, errors]) -> bytes\n\
@@ -275,23 +275,39 @@ impl PyBytes {
275275
#[pymethod(name = "endswith")]
276276
fn endswith(
277277
&self,
278-
suffix: Either<PyByteInner, PyTupleRef>,
279-
start: OptionalArg<PyObjectRef>,
280-
end: OptionalArg<PyObjectRef>,
278+
suffix: PyObjectRef,
279+
start: OptionalArg<Option<isize>>,
280+
end: OptionalArg<Option<isize>>,
281281
vm: &VirtualMachine,
282282
) -> PyResult<bool> {
283-
self.inner.startsendswith(suffix, start, end, true, vm)
283+
self.inner.elements[..].py_startsendswith(
284+
suffix,
285+
start,
286+
end,
287+
"endswith",
288+
"bytes",
289+
|s, x: &PyByteInner| s.ends_with(&x.elements[..]),
290+
vm,
291+
)
284292
}
285293

286294
#[pymethod(name = "startswith")]
287295
fn startswith(
288296
&self,
289-
prefix: Either<PyByteInner, PyTupleRef>,
290-
start: OptionalArg<PyObjectRef>,
291-
end: OptionalArg<PyObjectRef>,
297+
prefix: PyObjectRef,
298+
start: OptionalArg<Option<isize>>,
299+
end: OptionalArg<Option<isize>>,
292300
vm: &VirtualMachine,
293301
) -> PyResult<bool> {
294-
self.inner.startsendswith(prefix, start, end, false, vm)
302+
self.inner.elements[..].py_startsendswith(
303+
prefix,
304+
start,
305+
end,
306+
"startswith",
307+
"bytes",
308+
|s, x: &PyByteInner| s.starts_with(&x.elements[..]),
309+
vm,
310+
)
295311
}
296312

297313
#[pymethod(name = "find")]

‎vm/src/obj/objstr.rs

Lines changed: 28 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,13 @@ use super::objsequence::PySliceableSequence;
2121
use super::objslice::PySliceRef;
2222
use super::objtuple;
2323
use super::objtype::{self, PyClassRef};
24-
use super::pystr::PyCommonString;
24+
use super::pystr::{adjust_indices, PyCommonString, StringRange};
2525
use crate::cformat::{
2626
CFormatPart, CFormatPreconversor, CFormatQuantity, CFormatSpec, CFormatString, CFormatType,
2727
CNumberType,
2828
};
2929
use crate::format::{FormatParseError, FormatPart, FormatPreconversor, FormatSpec, FormatString};
30-
use crate::function::{single_or_tuple_any, OptionalArg, PyFuncArgs};
30+
use crate::function::{OptionalArg, PyFuncArgs};
3131
use crate::pyhash;
3232
use crate::pyobject::{
3333
Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyIterable,
@@ -525,23 +525,15 @@ impl PyString {
525525
end: OptionalArg<Option<isize>>,
526526
vm: &VirtualMachine,
527527
) -> PyResult<bool> {
528-
let range = adjust_indices(start, end, self.value.len());
529-
if range.is_normal() {
530-
let value = &self.value[range];
531-
single_or_tuple_any(
532-
suffix,
533-
|s: &PyStringRef| Ok(value.ends_with(&s.value)),
534-
|o| {
535-
format!(
536-
"endswith first arg must be str or a tuple of str, not {}",
537-
o.class(),
538-
)
539-
},
540-
vm,
541-
)
542-
} else {
543-
Ok(false)
544-
}
528+
self.value.as_str().py_startsendswith(
529+
suffix,
530+
start,
531+
end,
532+
"endswith",
533+
"str",
534+
|s, x: &PyStringRef| s.ends_with(x.as_str()),
535+
vm,
536+
)
545537
}
546538

547539
#[pymethod]
@@ -552,23 +544,15 @@ impl PyString {
552544
end: OptionalArg<Option<isize>>,
553545
vm: &VirtualMachine,
554546
) -> PyResult<bool> {
555-
let range = adjust_indices(start, end, self.value.len());
556-
if range.is_normal() {
557-
let value = &self.value[range];
558-
single_or_tuple_any(
559-
prefix,
560-
|s: &PyStringRef| Ok(value.starts_with(&s.value)),
561-
|o| {
562-
format!(
563-
"startswith first arg must be str or a tuple of str, not {}",
564-
o.class(),
565-
)
566-
},
567-
vm,
568-
)
569-
} else {
570-
Ok(false)
571-
}
547+
self.value.as_str().py_startsendswith(
548+
prefix,
549+
start,
550+
end,
551+
"startswith",
552+
"str",
553+
|s, x: &PyStringRef| s.starts_with(x.as_str()),
554+
vm,
555+
)
572556
}
573557

574558
#[pymethod]
@@ -1724,41 +1708,6 @@ impl PySliceableSequence for String {
17241708
}
17251709
}
17261710

1727-
pub trait StringRange {
1728-
fn is_normal(&self) -> bool;
1729-
}
1730-
1731-
impl StringRange for std::ops::Range<usize> {
1732-
fn is_normal(&self) -> bool {
1733-
self.start <= self.end
1734-
}
1735-
}
1736-
1737-
// help get optional string indices
1738-
pub fn adjust_indices(
1739-
start: OptionalArg<Option<isize>>,
1740-
end: OptionalArg<Option<isize>>,
1741-
len: usize,
1742-
) -> std::ops::Range<usize> {
1743-
let mut start = start.flat_option().unwrap_or(0);
1744-
let mut end = end.flat_option().unwrap_or(len as isize);
1745-
if end > len as isize {
1746-
end = len as isize;
1747-
} else if end < 0 {
1748-
end += len as isize;
1749-
if end < 0 {
1750-
end = 0;
1751-
}
1752-
}
1753-
if start < 0 {
1754-
start += len as isize;
1755-
if start < 0 {
1756-
start = 0;
1757-
}
1758-
}
1759-
start as usize..end as usize
1760-
}
1761-
17621711
// According to python following categories aren't printable:
17631712
// * Cc (Other, Control)
17641713
// * Cf (Other, Format)
@@ -1851,6 +1800,14 @@ mod tests {
18511800
}
18521801

18531802
impl PyCommonString<'_, char> for str {
1803+
fn get_slice(&self, range: std::ops::Range<usize>) -> &Self {
1804+
&self[range]
1805+
}
1806+
1807+
fn len(&self) -> usize {
1808+
Self::len(self)
1809+
}
1810+
18541811
fn py_split_whitespace<F>(&self, maxsplit: isize, convert: F) -> Vec<PyObjectRef>
18551812
where
18561813
F: Fn(&Self) -> PyObjectRef,

0 commit comments

Comments
 (0)