Skip to content

Commit 8a7c46d

Browse files
jgirardetjgirardet
committed
add endswith startswith
1 parent 14658b6 commit 8a7c46d

4 files changed

Lines changed: 116 additions & 2 deletions

File tree

‎tests/snippets/bytes.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@
162162
with assertRaises(TypeError):
163163
b"b".center(2, b"ba")
164164
b"kok".center(5, bytearray(b"x"))
165-
b"kok".center(-5,)
165+
b"kok".center(-5)
166166

167167
# count
168168
assert b"azeazerazeazopia".count(b"aze") == 3
@@ -191,3 +191,19 @@
191191
)
192192
with assertRaises(TypeError):
193193
b"".join((b"km", "kl"))
194+
195+
196+
# endswith startswith
197+
assert b"abcde".endswith(b"de")
198+
assert b"abcde".endswith(b"")
199+
assert not b"abcde".endswith(b"zx")
200+
assert b"abcde".endswith(b"bc", 0, 3)
201+
assert not b"abcde".endswith(b"bc", 2, 3)
202+
assert b"abcde".endswith((b"c", b"de"))
203+
204+
assert b"abcde".startswith(b"ab")
205+
assert b"abcde".startswith(b"")
206+
assert not b"abcde".startswith(b"zx")
207+
assert b"abcde".startswith(b"cd", 2)
208+
assert not b"abcde".startswith(b"cd", 1, 4)
209+
assert b"abcde".startswith((b"a", b"bc"))

‎vm/src/obj/objbyteinner.rs‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ use std::collections::hash_map::DefaultHasher;
1212
use std::hash::{Hash, Hasher};
1313

1414
use super::objint;
15-
use super::objsequence::PySliceableSequence;
15+
use super::objtype;
16+
use super::objsequence::{PySliceableSequence, is_valid_slice_arg};
1617
use crate::obj::objint::PyInt;
1718
use num_integer::Integer;
1819
use num_traits::ToPrimitive;
@@ -21,6 +22,7 @@ use super::objbytearray::{get_value as get_value_bytearray, PyByteArray};
2122
use super::objbytes::PyBytes;
2223
use super::objmemory::PyMemoryView;
2324
use super::objnone::PyNone;
25+
use super::objsequence;
2426

2527
#[derive(Debug, Default, Clone)]
2628
pub struct PyByteInner {
@@ -495,6 +497,61 @@ impl PyByteInner {
495497

496498
Ok(vm.ctx.new_bytes(refs))
497499
}
500+
501+
pub fn startsendswith(
502+
&self,
503+
arg: PyObjectRef,
504+
start: OptionalArg<PyObjectRef>,
505+
end: OptionalArg<PyObjectRef>,
506+
endswith: bool, // true for endswith, false for startswith
507+
vm: &VirtualMachine,
508+
) -> PyResult {
509+
let suff = if objtype::isinstance(&arg, &vm.ctx.tuple_type()) {
510+
let mut flatten = vec![];
511+
for v in objsequence::get_elements(&arg).to_vec() {
512+
match try_as_bytes_like(&v) {
513+
None => {
514+
return Err(vm.new_type_error(format!(
515+
"a bytes-like object is required, not {}",
516+
&v.class().name,
517+
)));
518+
}
519+
Some(value) => flatten.extend(value),
520+
}
521+
}
522+
flatten
523+
} else {
524+
match try_as_bytes_like(&arg) {
525+
Some(value) => value,
526+
None => {
527+
return Err(vm.new_type_error(format!(
528+
"endswith first arg must be bytes or a tuple of bytes, not {}",
529+
arg
530+
)));
531+
}
532+
}
533+
};
534+
535+
if suff.is_empty() {
536+
return Ok(vm.new_bool(true));
537+
}
538+
let range = self.elements.get_slice_range(
539+
&is_valid_slice_arg(start, vm)?,
540+
&is_valid_slice_arg(end, vm)?,
541+
);
542+
543+
if range.end - range.start < suff.len() {
544+
return Ok(vm.new_bool(false));
545+
}
546+
547+
let offset = if endswith {
548+
(range.end - suff.len())..range.end
549+
} else {
550+
0..suff.len()
551+
};
552+
553+
Ok(vm.new_bool(suff.as_slice() == &self.elements.do_slice(range)[offset]))
554+
}
498555
}
499556

500557
pub fn try_as_byte(obj: &PyObjectRef) -> Option<Vec<u8>> {

‎vm/src/obj/objbytes.rs‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,28 @@ impl PyBytesRef {
258258
fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult {
259259
self.inner.join(iter, vm)
260260
}
261+
262+
#[pymethod(name = "endswith")]
263+
fn endswith(
264+
self,
265+
suffix: PyObjectRef,
266+
start: OptionalArg<PyObjectRef>,
267+
end: OptionalArg<PyObjectRef>,
268+
vm: &VirtualMachine,
269+
) -> PyResult {
270+
self.inner.startsendswith(suffix, start, end, true, vm)
271+
}
272+
273+
#[pymethod(name = "startswith")]
274+
fn startswith(
275+
self,
276+
suffix: PyObjectRef,
277+
start: OptionalArg<PyObjectRef>,
278+
end: OptionalArg<PyObjectRef>,
279+
vm: &VirtualMachine,
280+
) -> PyResult {
281+
self.inner.startsendswith(suffix, start, end, false, vm)
282+
}
261283
}
262284

263285
#[derive(Debug)]

‎vm/src/obj/objsequence.rs‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
use crate::function::OptionalArg;
2+
use crate::obj::objnone::PyNone;
13
use std::cell::RefCell;
24
use std::marker::Sized;
35
use std::ops::{Deref, DerefMut, Range};
@@ -371,3 +373,20 @@ pub fn get_mut_elements<'a>(obj: &'a PyObjectRef) -> impl DerefMut<Target = Vec<
371373
}
372374
panic!("Cannot extract elements from non-sequence");
373375
}
376+
377+
//Check if given arg could be used with PySciceableSequance.get_slice_range()
378+
pub fn is_valid_slice_arg(
379+
arg: OptionalArg<PyObjectRef>,
380+
vm: &VirtualMachine,
381+
) -> Result<Option<BigInt>, PyObjectRef> {
382+
if let OptionalArg::Present(value) = arg {
383+
match_class!(value,
384+
i @ PyInt => Ok(Some(i.as_bigint().clone())),
385+
_obj @ PyNone => Ok(None),
386+
_=> {return Err(vm.new_type_error("slice indices must be integers or None or have an __index__ method".to_string()));}
387+
// TODO: check for an __index__ method
388+
)
389+
} else {
390+
Ok(None)
391+
}
392+
}

0 commit comments

Comments
 (0)