Skip to content

Commit e19cb3e

Browse files
youknowoneyouknowone
authored
Merge pull request #1855 from youknowone/strbytes-partition
Fix str, bytes partition
2 parents f2fbb07 + 7d7d048 commit e19cb3e

File tree

6 files changed

+101
-63
lines changed

6 files changed

+101
-63
lines changed

‎Lib/test/test_bytes.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -737,13 +737,11 @@ def test_rsplit_unicodewhitespace(self):
737737
b = self.type2test(b"\x09\x0A\x0B\x0C\x0D\x1C\x1D\x1E\x1F")
738738
self.assertEqual(b.rsplit(), [b'\x1c\x1d\x1e\x1f'])
739739

740-
@unittest.skip("TODO: RUSTPYTHON")
741740
def test_partition(self):
742741
b = self.type2test(b'mississippi')
743742
self.assertEqual(b.partition(b'ss'), (b'mi', b'ss', b'issippi'))
744743
self.assertEqual(b.partition(b'w'), (b'mississippi', b'', b''))
745744

746-
@unittest.skip("TODO: RUSTPYTHON")
747745
def test_rpartition(self):
748746
b = self.type2test(b'mississippi')
749747
self.assertEqual(b.rpartition(b'ss'), (b'missi', b'ss', b'ippi'))
@@ -1578,7 +1576,6 @@ def test_copied(self):
15781576
x = bytearray(b'')
15791577
self.assertIsNot(x, x.translate(t))
15801578

1581-
@unittest.skip("TODO: RUSTPYTHON")
15821579
def test_partition_bytearray_doesnt_share_nullstring(self):
15831580
a, b, c = bytearray(b"x").partition(b"y")
15841581
self.assertEqual(b, b"")

‎Lib/test/test_unicode.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -419,8 +419,6 @@ def test_rsplit(self):
419419
self.checkequal([left, right],
420420
left + delim * 2 + right, 'rsplit', delim *2)
421421

422-
# TODO: RUSTPYTHON
423-
@unittest.expectedFailure
424422
def test_partition(self):
425423
string_tests.MixinStrUnicodeUserStringTest.test_partition(self)
426424
# test mixed kinds
@@ -438,8 +436,6 @@ def test_partition(self):
438436
self.checkequal((left, delim * 2, right),
439437
left + delim * 2 + right, 'partition', delim * 2)
440438

441-
# TODO: RUSTPYTHON
442-
@unittest.expectedFailure
443439
def test_rpartition(self):
444440
string_tests.MixinStrUnicodeUserStringTest.test_rpartition(self)
445441
# test mixed kinds

‎vm/src/obj/objbytearray.rs

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -413,21 +413,25 @@ impl PyByteArray {
413413
fn partition(&self, sep: PyByteInner, vm: &VirtualMachine) -> PyResult {
414414
// sep ALWAYS converted to bytearray even it's bytes or memoryview
415415
// so its ok to accept PyByteInner
416-
let (left, right) = self.borrow_value().partition(&sep, false)?;
416+
let value = self.borrow_value();
417+
let (front, has_mid, back) = value.partition(&sep, vm)?;
417418
Ok(vm.ctx.new_tuple(vec![
418-
vm.ctx.new_bytearray(left),
419-
vm.ctx.new_bytearray(sep.elements),
420-
vm.ctx.new_bytearray(right),
419+
vm.ctx.new_bytearray(front.to_vec()),
420+
vm.ctx
421+
.new_bytearray(if has_mid { sep.elements } else { Vec::new() }),
422+
vm.ctx.new_bytearray(back.to_vec()),
421423
]))
422424
}
423425

424426
#[pymethod(name = "rpartition")]
425427
fn rpartition(&self, sep: PyByteInner, vm: &VirtualMachine) -> PyResult {
426-
let (left, right) = self.borrow_value().partition(&sep, true)?;
428+
let value = self.borrow_value();
429+
let (front, has_mid, back) = value.rpartition(&sep, vm)?;
427430
Ok(vm.ctx.new_tuple(vec![
428-
vm.ctx.new_bytearray(left),
429-
vm.ctx.new_bytearray(sep.elements),
430-
vm.ctx.new_bytearray(right),
431+
vm.ctx.new_bytearray(front.to_vec()),
432+
vm.ctx
433+
.new_bytearray(if has_mid { sep.elements } else { Vec::new() }),
434+
vm.ctx.new_bytearray(back.to_vec()),
431435
]))
432436
}
433437

‎vm/src/obj/objbyteinner.rs

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1012,13 +1012,42 @@ impl PyByteInner {
10121012
}
10131013
}
10141014

1015-
pub fn partition(&self, sep: &PyByteInner, reverse: bool) -> PyResult<(Vec<u8>, Vec<u8>)> {
1016-
let splitted = if reverse {
1017-
split_slice_reverse(&self.elements, &sep.elements, 1)
1015+
pub fn partition(
1016+
&self,
1017+
sub: &PyByteInner,
1018+
vm: &VirtualMachine,
1019+
) -> PyResult<(Vec<u8>, bool, Vec<u8>)> {
1020+
if sub.elements.is_empty() {
1021+
return Err(vm.new_value_error("empty separator".to_owned()));
1022+
}
1023+
1024+
let mut sp = self.elements.splitn_str(2, &sub.elements);
1025+
let front = sp.next().unwrap().to_vec();
1026+
let (has_mid, back) = if let Some(back) = sp.next() {
1027+
(true, back.to_vec())
1028+
} else {
1029+
(false, Vec::new())
1030+
};
1031+
Ok((front, has_mid, back))
1032+
}
1033+
1034+
pub fn rpartition(
1035+
&self,
1036+
sub: &PyByteInner,
1037+
vm: &VirtualMachine,
1038+
) -> PyResult<(Vec<u8>, bool, Vec<u8>)> {
1039+
if sub.elements.is_empty() {
1040+
return Err(vm.new_value_error("empty separator".to_owned()));
1041+
}
1042+
1043+
let mut sp = self.elements.rsplitn_str(2, &sub.elements);
1044+
let back = sp.next().unwrap().to_vec();
1045+
let (has_mid, front) = if let Some(front) = sp.next() {
1046+
(true, front.to_vec())
10181047
} else {
1019-
split_slice(&self.elements, &sep.elements, 1)
1048+
(false, Vec::new())
10201049
};
1021-
Ok((splitted[0].to_vec(), splitted[1].to_vec()))
1050+
Ok((front, has_mid, back))
10221051
}
10231052

10241053
pub fn expandtabs(&self, options: ByteInnerExpandtabsOptions) -> Vec<u8> {

‎vm/src/obj/objbytes.rs

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -366,21 +366,32 @@ impl PyBytes {
366366

367367
#[pymethod(name = "partition")]
368368
fn partition(&self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
369-
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
370-
371-
let (left, right) = self.inner.partition(&sepa, false)?;
372-
Ok(vm
373-
.ctx
374-
.new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
369+
let sub = PyByteInner::try_from_object(vm, sep.clone())?;
370+
let (front, has_mid, back) = self.inner.partition(&sub, vm)?;
371+
Ok(vm.ctx.new_tuple(vec![
372+
vm.ctx.new_bytes(front),
373+
if has_mid {
374+
sep
375+
} else {
376+
vm.ctx.new_bytes(Vec::new())
377+
},
378+
vm.ctx.new_bytes(back),
379+
]))
375380
}
381+
376382
#[pymethod(name = "rpartition")]
377383
fn rpartition(&self, sep: PyObjectRef, vm: &VirtualMachine) -> PyResult {
378-
let sepa = PyByteInner::try_from_object(vm, sep.clone())?;
379-
380-
let (left, right) = self.inner.partition(&sepa, true)?;
381-
Ok(vm
382-
.ctx
383-
.new_tuple(vec![vm.ctx.new_bytes(left), sep, vm.ctx.new_bytes(right)]))
384+
let sub = PyByteInner::try_from_object(vm, sep.clone())?;
385+
let (front, has_mid, back) = self.inner.rpartition(&sub, vm)?;
386+
Ok(vm.ctx.new_tuple(vec![
387+
vm.ctx.new_bytes(front),
388+
if has_mid {
389+
sep
390+
} else {
391+
vm.ctx.new_bytes(Vec::new())
392+
},
393+
vm.ctx.new_bytes(back),
394+
]))
384395
}
385396

386397
#[pymethod(name = "expandtabs")]

‎vm/src/obj/objstr.rs

Lines changed: 32 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -954,42 +954,43 @@ impl PyString {
954954
}
955955

956956
#[pymethod]
957-
fn partition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyObjectRef {
958-
let value = &self.value;
959-
let sub = &sub.value;
960-
let mut new_tup = Vec::new();
961-
if value.contains(sub) {
962-
new_tup = value
963-
.splitn(2, sub)
964-
.map(|s| vm.ctx.new_str(s.to_owned()))
965-
.collect();
966-
new_tup.insert(1, vm.ctx.new_str(sub.clone()));
967-
} else {
968-
new_tup.push(vm.ctx.new_str(value.clone()));
969-
new_tup.push(vm.ctx.new_str("".to_owned()));
970-
new_tup.push(vm.ctx.new_str("".to_owned()));
957+
fn partition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyResult {
958+
if sub.value.is_empty() {
959+
return Err(vm.new_value_error("empty separator".to_owned()));
971960
}
972-
vm.ctx.new_tuple(new_tup)
961+
let mut sp = self.value.splitn(2, &sub.value);
962+
let front = sp.next().unwrap();
963+
let elems = if let Some(back) = sp.next() {
964+
[front, &sub.value, back]
965+
} else {
966+
[front, "", ""]
967+
};
968+
Ok(vm.ctx.new_tuple(
969+
elems
970+
.iter()
971+
.map(|&s| vm.ctx.new_str(s.to_owned()))
972+
.collect(),
973+
))
973974
}
974975

975976
#[pymethod]
976-
fn rpartition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyObjectRef {
977-
let value = &self.value;
978-
let sub = &sub.value;
979-
let mut new_tup = Vec::new();
980-
if value.contains(sub) {
981-
new_tup = value
982-
.rsplitn(2, sub)
983-
.map(|s| vm.ctx.new_str(s.to_owned()))
984-
.collect();
985-
new_tup.swap(0, 1); // so it's in the right order
986-
new_tup.insert(1, vm.ctx.new_str(sub.clone()));
987-
} else {
988-
new_tup.push(vm.ctx.new_str("".to_owned()));
989-
new_tup.push(vm.ctx.new_str("".to_owned()));
990-
new_tup.push(vm.ctx.new_str(value.clone()));
977+
fn rpartition(&self, sub: PyStringRef, vm: &VirtualMachine) -> PyResult {
978+
if sub.value.is_empty() {
979+
return Err(vm.new_value_error("empty separator".to_owned()));
991980
}
992-
vm.ctx.new_tuple(new_tup)
981+
let mut sp = self.value.rsplitn(2, &sub.value);
982+
let back = sp.next().unwrap();
983+
let elems = if let Some(front) = sp.next() {
984+
[front, &sub.value, back]
985+
} else {
986+
["", "", back]
987+
};
988+
Ok(vm.ctx.new_tuple(
989+
elems
990+
.iter()
991+
.map(|&s| vm.ctx.new_str(s.to_owned()))
992+
.collect(),
993+
))
993994
}
994995

995996
/// Return `true` if the sequence is ASCII titlecase and the sequence is not

0 commit comments

Comments
 (0)