Skip to content

Commit ed94ddb

Browse files
committed
add split, rsplit, partition, rpartition, expandtabs, spitlines, zfill,
replace to bytearray
1 parent abf0a32 commit ed94ddb

File tree

3 files changed

+322
-7
lines changed

3 files changed

+322
-7
lines changed

tests/snippets/bytearray.py

Lines changed: 237 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,243 @@
357357
assert bytearray(b"mississippi").rstrip(b"ipz") == bytearray(b"mississ")
358358

359359

360+
361+
# split
362+
assert bytearray(b"1,2,3").split(bytearray(b",")) == [bytearray(b"1"), bytearray(b"2"), bytearray(b"3")]
363+
assert bytearray(b"1,2,3").split(bytearray(b","), maxsplit=1) == [bytearray(b"1"), bytearray(b"2,3")]
364+
assert bytearray(b"1,2,,3,").split(bytearray(b",")) == [bytearray(b"1"), bytearray(b"2"), bytearray(b""), bytearray(b"3"), bytearray(b"")]
365+
assert bytearray(b"1 2 3").split() == [bytearray(b"1"), bytearray(b"2"), bytearray(b"3")]
366+
assert bytearray(b"1 2 3").split(maxsplit=1) == [bytearray(b"1"), bytearray(b"2 3")]
367+
assert bytearray(b" 1 2 3 ").split() == [bytearray(b"1"), bytearray(b"2"), bytearray(b"3")]
368+
assert bytearray(b"k\ruh\nfz e f").split() == [bytearray(b"k"), bytearray(b"uh"), bytearray(b"fz"), bytearray(b"e"), bytearray(b"f")]
369+
assert bytearray(b"Two lines\n").split(bytearray(b"\n")) == [bytearray(b"Two lines"), bytearray(b"")]
370+
assert bytearray(b"").split() == []
371+
assert bytearray(b"").split(bytearray(b"\n")) == [bytearray(b"")]
372+
assert bytearray(b"\n").split(bytearray(b"\n")) == [bytearray(b""), bytearray(b"")]
373+
374+
SPLIT_FIXTURES = [
375+
[
376+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
377+
[4, 5],
378+
[[1, 2, 3], [1, 2, 3], [1, 2, 3]],
379+
[[1, 2, 3], [1, 2, 3], [1, 2, 3]],
380+
-1,
381+
],
382+
[
383+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5],
384+
[4, 5],
385+
[[1, 2, 3], [1, 2, 3], [1, 2, 3], []],
386+
[[1, 2, 3], [1, 2, 3], [1, 2, 3], []],
387+
-1,
388+
],
389+
[
390+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 3],
391+
[4, 5],
392+
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [3]],
393+
[[1, 2, 3], [1, 2, 3], [1, 2, 3], [3]],
394+
-1,
395+
],
396+
[
397+
[4, 5, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
398+
[4, 5],
399+
[[], [2, 3], [1, 2, 3], [1, 2, 3]],
400+
[[], [2, 3], [1, 2, 3], [1, 2, 3]],
401+
-1,
402+
],
403+
[
404+
[1, 4, 5, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
405+
[4, 5],
406+
[[1], [2, 3], [1, 2, 3], [1, 2, 3]],
407+
[[1], [2, 3], [1, 2, 3], [1, 2, 3]],
408+
-1,
409+
],
410+
[
411+
[1, 2, 3, 4, 5, 4, 5, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
412+
[4, 5],
413+
[[1, 2, 3], [], [], [1, 2, 3], [1, 2, 3]],
414+
[[1, 2, 3], [], [], [1, 2, 3], [1, 2, 3]],
415+
-1,
416+
],
417+
# maxsplit
418+
[
419+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
420+
[4, 5],
421+
[[1, 2, 3], [1, 2, 3, 4, 5, 1, 2, 3]],
422+
[[1, 2, 3, 4, 5, 1, 2, 3], [1, 2, 3]],
423+
1,
424+
],
425+
[
426+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5],
427+
[4, 5],
428+
[[1, 2, 3], [1, 2, 3, 4, 5, 1, 2, 3, 4, 5]],
429+
[[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3], []],
430+
1,
431+
],
432+
[
433+
[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 3],
434+
[4, 5],
435+
[[1, 2, 3], [1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 3]],
436+
[[1, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3], [3]],
437+
1,
438+
],
439+
[
440+
[4, 5, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
441+
[4, 5],
442+
[[], [2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3]],
443+
[[4, 5, 2, 3, 4, 5, 1, 2, 3], [1, 2, 3]],
444+
1,
445+
],
446+
[
447+
[1, 4, 5, 2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
448+
[4, 5],
449+
[[1], [2, 3, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3]],
450+
[[1, 4, 5, 2, 3, 4, 5, 1, 2, 3], [1, 2, 3]],
451+
1,
452+
],
453+
[
454+
[1, 2, 3, 4, 5, 4, 5, 4, 5, 1, 2, 3, 4, 5, 1, 2, 3],
455+
[4, 5],
456+
[[1, 2, 3], [], [4, 5, 1, 2, 3, 4, 5, 1, 2, 3]],
457+
[[1, 2, 3, 4, 5, 4, 5], [1, 2, 3], [1, 2, 3]],
458+
2,
459+
],
460+
[
461+
[13, 13, 13, 117, 104, 10, 102, 122, 32, 101, 102, 9, 9],
462+
None,
463+
[[117, 104], [102, 122], [101, 102]],
464+
[[117, 104], [102, 122], [101, 102]],
465+
-1,
466+
],
467+
[
468+
[13, 13, 13, 117, 104, 10, 102, 122, 32, 101, 102, 9, 9],
469+
None,
470+
[[117, 104, 10, 102, 122, 32, 101, 102, 9, 9]],
471+
[[13, 13, 13, 117, 104, 10, 102, 122, 32, 101, 102]],
472+
0,
473+
],
474+
[
475+
[13, 13, 13, 117, 104, 10, 102, 122, 32, 101, 102, 9, 9],
476+
None,
477+
[[117, 104], [102, 122, 32, 101, 102, 9, 9]],
478+
[[13, 13, 13, 117, 104, 10, 102, 122], [101, 102]],
479+
1,
480+
],
481+
[
482+
[13, 13, 13, 117, 104, 10, 102, 122, 32, 101, 102, 9, 9],
483+
None,
484+
[[117, 104], [102, 122], [101, 102, 9, 9]],
485+
[[13, 13, 13, 117, 104], [102, 122], [101, 102]],
486+
2,
487+
],
488+
[
489+
[13, 13, 13, 117, 104, 10, 10, 10, 102, 122, 32, 32, 101, 102, 9, 9],
490+
None,
491+
[[117, 104], [102, 122], [101, 102]],
492+
[[117, 104], [102, 122], [101, 102]],
493+
-1,
494+
],
495+
[[49, 44, 50, 44, 51], [44], [[49], [50], [51]], [[49], [50], [51]], -1],
496+
[[49, 44, 50, 44, 51], [44], [[49], [50, 44, 51]], [[49, 44, 50], [51]], 1],
497+
[
498+
[49, 44, 50, 44, 44, 51, 44],
499+
[44],
500+
[[49], [50], [], [51], []],
501+
[[49], [50], [], [51], []],
502+
-1,
503+
],
504+
[[49, 32, 50, 32, 51], None, [[49], [50], [51]], [[49], [50], [51]], -1],
505+
[[49, 32, 50, 32, 51], None, [[49], [50, 32, 51]], [[49, 32, 50], [51]], 1],
506+
[
507+
[32, 32, 32, 49, 32, 32, 32, 50, 32, 32, 32, 51, 32, 32, 32],
508+
None,
509+
[[49], [50], [51]],
510+
[[49], [50], [51]],
511+
-1,
512+
],
513+
]
514+
515+
516+
# for i in SPLIT_FIXTURES: # for not yet implemented : TypeError: Unsupported method: __next__
517+
n_sp = 0
518+
while n_sp < len(SPLIT_FIXTURES):
519+
i = SPLIT_FIXTURES[n_sp]
520+
sep = None if i[1] == None else bytearray(i[1])
521+
try:
522+
assert bytearray(i[0]).split(sep=sep, maxsplit=i[4]) == [bytearray(j) for j in i[2]]
523+
except AssertionError:
524+
print(i[0], i[1], i[2])
525+
print(
526+
"Expected : ", [list(x) for x in bytearray(i[0]).split(sep=sep, maxsplit=i[4])]
527+
)
528+
break
529+
530+
try:
531+
assert bytearray(i[0]).rsplit(sep=sep, maxsplit=i[4]) == [bytearray(j) for j in i[3]]
532+
except AssertionError:
533+
print(i[0], i[1], i[2])
534+
print(
535+
"Expected Rev : ",
536+
[list(x) for x in bytearray(i[0]).rsplit(sep=sep, maxsplit=i[4])],
537+
)
538+
break
539+
540+
n_sp += 1
541+
542+
543+
# expandtabs
544+
a = bytearray(b"\x01\x03\r\x05\t8CYZ\t\x06CYZ\t\x17cba`\n\x12\x13\x14")
545+
assert (
546+
a.expandtabs() == bytearray(b"\x01\x03\r\x05 8CYZ \x06CYZ \x17cba`\n\x12\x13\x14")
547+
)
548+
assert a.expandtabs(5) == bytearray(b"\x01\x03\r\x05 8CYZ \x06CYZ \x17cba`\n\x12\x13\x14")
549+
assert bytearray(b"01\t012\t0123\t01234").expandtabs() == bytearray(b"01 012 0123 01234")
550+
assert bytearray(b"01\t012\t0123\t01234").expandtabs(4) == bytearray(b"01 012 0123 01234")
551+
assert bytearray(b"123\t123").expandtabs(-5) == bytearray(b"123123")
552+
assert bytearray(b"123\t123").expandtabs(0) == bytearray(b"123123")
553+
554+
555+
# # partition
556+
assert bytearray(b"123456789").partition(b"45") == ((b"123"), bytearray(b"45"), bytearray(b"6789"))
557+
assert bytearray(b"14523456789").partition(b"45") == ((b"1"), bytearray(b"45"), bytearray(b"23456789"))
558+
a = bytearray(b"14523456789").partition(b"45")
559+
assert isinstance(a[1], bytearray)
560+
a = bytearray(b"14523456789").partition(memoryview(b"45"))
561+
assert isinstance(a[1], bytearray)
562+
563+
# partition
564+
assert bytearray(b"123456789").rpartition(bytearray(b"45")) == ((bytearray(b"123")), bytearray(b"45"), bytearray(b"6789"))
565+
assert bytearray(b"14523456789").rpartition(bytearray(b"45")) == ((bytearray(b"14523")), bytearray(b"45"), bytearray(b"6789"))
566+
a = bytearray(b"14523456789").rpartition(b"45")
567+
assert isinstance(a[1], bytearray)
568+
a = bytearray(b"14523456789").rpartition(memoryview(b"45"))
569+
assert isinstance(a[1], bytearray)
570+
571+
# splitlines
572+
assert bytearray(b"ab c\n\nde fg\rkl\r\n").splitlines() == [bytearray(b"ab c"), bytearray(b""), bytearray(b"de fg"), bytearray(b"kl")]
573+
assert bytearray(b"ab c\n\nde fg\rkl\r\n").splitlines(keepends=True) == [
574+
bytearray(b"ab c\n"),
575+
bytearray(b"\n"),
576+
bytearray(b"de fg\r"),
577+
bytearray(b"kl\r\n"),
578+
]
579+
assert bytearray(b"").splitlines() == []
580+
assert bytearray(b"One line\n").splitlines() == [b"One line"]
581+
582+
# zfill
583+
584+
assert bytearray(b"42").zfill(5) == bytearray(b"00042")
585+
assert bytearray(b"-42").zfill(5) == bytearray(b"-0042")
586+
assert bytearray(b"42").zfill(1) == bytearray(b"42")
587+
assert bytearray(b"42").zfill(-1) == bytearray(b"42")
588+
589+
# replace
590+
assert bytearray(b"123456789123").replace(b"23",b"XX") ==bytearray(b'1XX4567891XX')
591+
assert bytearray(b"123456789123").replace(b"23",b"XX", 1) ==bytearray(b'1XX456789123')
592+
assert bytearray(b"123456789123").replace(b"23",b"XX", 0) == bytearray(b"123456789123")
593+
assert bytearray(b"123456789123").replace(b"23",b"XX", -1) ==bytearray(b'1XX4567891XX')
594+
assert bytearray(b"123456789123").replace(b"23", bytearray(b"")) == bytearray(b"14567891")
595+
596+
360597
# clear
361598
a = bytearray(b"abcd")
362599
a.clear()
@@ -387,5 +624,3 @@
387624
assert a == bytearray(b"append")
388625
assert len(a) == 6
389626
assert a.pop() == 100
390-
391-
import bytes as bbytes

vm/src/obj/objbytearray.rs

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
33
use crate::function::OptionalArg;
44
use crate::obj::objbyteinner::{
5-
ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerPosition,
5+
ByteInnerExpandtabsOptions, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
6+
ByteInnerPosition, ByteInnerSplitOptions, ByteInnerSplitlinesOptions,
67
ByteInnerTranslateOptions, ByteOr, PyByteInner,
78
};
89
use crate::obj::objint::PyIntRef;
@@ -351,6 +352,89 @@ impl PyByteArrayRef {
351352
))
352353
}
353354

355+
#[pymethod(name = "split")]
356+
fn split(self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
357+
let as_bytes = self
358+
.inner
359+
.borrow()
360+
.split(options, false)?
361+
.iter()
362+
.map(|x| vm.ctx.new_bytearray(x.to_vec()))
363+
.collect::<Vec<PyObjectRef>>();
364+
Ok(vm.ctx.new_list(as_bytes))
365+
}
366+
367+
#[pymethod(name = "rsplit")]
368+
fn rsplit(self, options: ByteInnerSplitOptions, vm: &VirtualMachine) -> PyResult {
369+
let as_bytes = self
370+
.inner
371+
.borrow()
372+
.split(options, true)?
373+
.iter()
374+
.map(|x| vm.ctx.new_bytearray(x.to_vec()))
375+
.collect::<Vec<PyObjectRef>>();
376+
Ok(vm.ctx.new_list(as_bytes))
377+
}
378+
379+
#[pymethod(name = "partition")]
380+
fn partition(self, sep: PyByteInner, vm: &VirtualMachine) -> PyResult {
381+
// sep ALWAYS converted to bytearray even it's bytes or memoryview
382+
// so its ok to accept PyByteInner
383+
let (left, right) = self.inner.borrow().partition(&sep, false)?;
384+
Ok(vm.ctx.new_tuple(vec![
385+
vm.ctx.new_bytearray(left),
386+
vm.ctx.new_bytearray(sep.elements),
387+
vm.ctx.new_bytearray(right),
388+
]))
389+
}
390+
391+
#[pymethod(name = "rpartition")]
392+
fn rpartition(self, sep: PyByteInner, vm: &VirtualMachine) -> PyResult {
393+
let (left, right) = self.inner.borrow().partition(&sep, true)?;
394+
Ok(vm.ctx.new_tuple(vec![
395+
vm.ctx.new_bytearray(left),
396+
vm.ctx.new_bytearray(sep.elements),
397+
vm.ctx.new_bytearray(right),
398+
]))
399+
}
400+
401+
#[pymethod(name = "expandtabs")]
402+
fn expandtabs(self, options: ByteInnerExpandtabsOptions, vm: &VirtualMachine) -> PyResult {
403+
Ok(vm
404+
.ctx
405+
.new_bytearray(self.inner.borrow().expandtabs(options)))
406+
}
407+
408+
#[pymethod(name = "splitlines")]
409+
fn splitlines(self, options: ByteInnerSplitlinesOptions, vm: &VirtualMachine) -> PyResult {
410+
let as_bytes = self
411+
.inner
412+
.borrow()
413+
.splitlines(options)
414+
.iter()
415+
.map(|x| vm.ctx.new_bytearray(x.to_vec()))
416+
.collect::<Vec<PyObjectRef>>();
417+
Ok(vm.ctx.new_list(as_bytes))
418+
}
419+
420+
#[pymethod(name = "zfill")]
421+
fn zfill(self, width: PyIntRef, vm: &VirtualMachine) -> PyResult {
422+
Ok(vm.ctx.new_bytearray(self.inner.borrow().zfill(width)))
423+
}
424+
425+
#[pymethod(name = "replace")]
426+
fn replace(
427+
self,
428+
old: PyByteInner,
429+
new: PyByteInner,
430+
count: OptionalArg<PyIntRef>,
431+
vm: &VirtualMachine,
432+
) -> PyResult {
433+
Ok(vm
434+
.ctx
435+
.new_bytearray(self.inner.borrow().replace(old, new, count)?))
436+
}
437+
354438
#[pymethod(name = "clear")]
355439
fn clear(self, _vm: &VirtualMachine) {
356440
self.inner.borrow_mut().elements.clear();

vm/src/obj/objbytes.rs

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -350,8 +350,6 @@ impl PyBytesRef {
350350

351351
#[pymethod(name = "partition")]
352352
fn partition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
353-
// TODO: when implementing bytearray,remember sep ALWAYS converted to bytearray
354-
// even it's bytes or memoryview so PyByteInner wiil be ok for args
355353
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
356354

357355
let (left, right) = self.inner.partition(&sepa, false)?;
@@ -361,8 +359,6 @@ impl PyBytesRef {
361359
}
362360
#[pymethod(name = "rpartition")]
363361
fn rpartition(self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
364-
// TODO: when implementing bytearray,remember sep ALWAYS converted to bytearray
365-
// even it's bytes or memoryview so PyByteInner wiil be ok for args
366362
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
367363

368364
let (left, right) = self.inner.partition(&sepa, true)?;

0 commit comments

Comments
 (0)