Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 0 additions & 17 deletions crates/vm/src/anystr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -321,23 +321,6 @@ pub(crate) trait AnyStr {
self.py_pad(width - len, 0, fillchar)
}

fn py_join(
&self,
mut iter: impl core::iter::Iterator<Item = PyResult<impl AnyStrWrapper<Self> + TryFromObject>>,
) -> PyResult<Self::Container> {
let mut joined = if let Some(elem) = iter.next() {
elem?.as_ref().unwrap().to_container()
} else {
return Ok(Self::Container::new());
};
for elem in iter {
let elem = elem?;
joined.push_str(self);
joined.push_str(elem.as_ref().unwrap());
}
Ok(joined)
}

fn py_partition<'a, F, S>(
&'a self,
sub: &Self,
Expand Down
15 changes: 10 additions & 5 deletions crates/vm/src/builtins/bytearray.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use crate::{
bytes_inner::{
ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
},
class::PyClassImpl,
common::{
Expand Down Expand Up @@ -354,10 +354,15 @@ impl PyByteArray {

#[pymethod]
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
// Driving the iterable runs Python, which can reach this bytearray,
// so the separator is taken by value rather than left borrowed.
let separator = self.inner().clone();
Ok(separator.join(iterable_of_bytes, vm)?.into())
// Export before driving the iterable: resizing is forbidden, but
// same-size mutations must be reflected in the joined result.
self.exports.fetch_add(1, Ordering::Release);
scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
let separator_len = self.__len__();
let items = BytesJoin::new(iterable_of_bytes, vm)?;
Ok(items
.join(separator_len, || self.borrow_buf().into(), vm)?
.into())
}

#[pymethod]
Expand Down
12 changes: 9 additions & 3 deletions crates/vm/src/builtins/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use crate::{
bytes_inner::{
ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub,
ByteInnerTranslateOptions, DecodeArgs, PyBytesInner, bytes_decode,
ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
},
class::{PyClassDef, PyClassImpl},
common::{hash::PyHash, lock::PyMutex},
Expand Down Expand Up @@ -350,8 +350,14 @@ impl PyBytes {
}

#[pymethod]
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
Ok(self.inner.join(iterable_of_bytes, vm)?.into())
fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
let items = BytesJoin::new(iterable_of_bytes, vm)?;
if let Some(only) = items.single_bytes(vm) {
return Ok(only);
}
Ok(vm
.ctx
.new_bytes(items.join(self.as_bytes().len(), || self.as_bytes().into(), vm)?))
}

#[pymethod]
Expand Down
21 changes: 2 additions & 19 deletions crates/vm/src/builtins/str.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1214,25 +1214,8 @@ impl PyStr {
// and which has its own wording for what it cannot iterate.
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
obj?.downcast::<Self>().map_err(|obj| {
vm.new_type_error(format!(
"sequence item {i}: expected str instance, {} found",
obj.class().slot_name()
))
})
});
let joined = match iter.exactly_one() {
Ok(first) => {
let first = first?;
if first.as_object().class().is(vm.ctx.types.str_type) {
return Ok(first);
}
first.as_wtf8().to_owned()
}
Err(iter) => zelf.as_wtf8().py_join(iter)?,
};
Ok(vm.ctx.new_str(joined))
let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
Self::join_items(&zelf, &items, vm)
}

/// `join` over already-materialized items: checks them and sizes the result before copying.
Expand Down
141 changes: 119 additions & 22 deletions crates/vm/src/bytes_inner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,12 @@ use crate::{
VirtualMachine,
anystr::{self, AnyStr, AnyStrContainer, AnyStrWrapper},
builtins::{
PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyStr, PyStrRef,
pystr, pystr::PyUtf8StrRef,
PyBaseExceptionRef, PyByteArray, PyBytes, PyBytesRef, PyInt, PyIntRef, PyList, PyListRef,
PyStr, PyStrRef, PyTuple, PyTupleRef, pystr, pystr::PyUtf8StrRef,
},
cformat::cformat_bytes,
common::hash,
common::wtf8::is_py_ascii_whitespace,
common::{borrow::BorrowedValue, hash},
function::{ArgIterable, Either, OptionalArg, PyComparisonValue},
literal::escape::Escape,
protocol::{BufferFlags, PyBuffer},
Expand Down Expand Up @@ -40,6 +40,122 @@ impl From<Vec<u8>> for PyBytesInner {
}
}

// Keep exact lists live while acquiring buffers: an exporter can replace items
// or resize the list. No list lock may span a call to the buffer protocol.
pub(crate) struct BytesJoin {
items: Either<PyListRef, PyTupleRef>,
}

impl BytesJoin {
pub(crate) fn new(iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Self> {
let items = if let Some(list) = iterable.downcast_ref_if_exact::<PyList>(vm) {
Either::A(list.to_owned())
} else if let Some(tuple) = iterable.downcast_ref_if_exact::<PyTuple>(vm) {
Either::B(tuple.to_owned())
} else {
// PySequence_Fast exhausts the iterator before inspecting its items.
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
let items = iterable.iter_sized(vm)?.collect::<PyResult<Vec<_>>>()?;
Either::B(vm.ctx.new_tuple(items))
};
Ok(Self { items })
}

fn len(&self) -> usize {
match &self.items {
Either::A(list) => list.borrow_vec().len(),
Either::B(tuple) => tuple.as_slice().len(),
}
}

fn item(&self, i: usize) -> Option<PyObjectRef> {
match &self.items {
Either::A(list) => list.borrow_vec().get(i).cloned(),
Either::B(tuple) => tuple.as_slice().get(i).cloned(),
}
}

pub(crate) fn single_bytes(&self, vm: &VirtualMachine) -> Option<PyBytesRef> {
if self.len() != 1 {
return None;
}
self.item(0)?
.downcast_ref_if_exact::<PyBytes>(vm)
.map(ToOwned::to_owned)
}

// stringlib_bytes_join
pub(crate) fn join<'a>(
&self,
separator_len: usize,
mut separator: impl FnMut() -> BorrowedValue<'a, [u8]>,
vm: &VirtualMachine,
) -> PyResult<Vec<u8>> {
let count = self.len();
let mut objects = Vec::with_capacity(count);
let mut buffers = Vec::new();
let mut len = 0usize;
let overflow = || vm.new_overflow_error("join() result is too long");
let changed = || vm.new_runtime_error("sequence changed size during iteration");
for i in 0..count {
let obj = self.item(i).ok_or_else(changed)?;
let item_len = if let Some(bytes) = obj.downcast_ref_if_exact::<PyBytes>(vm) {
bytes.as_bytes().len()
} else {
// Whatever the buffer lookup ran into, the item is only ever
// reported as the wrong kind of thing.
let type_error = || {
vm.new_type_error(format!(
"sequence item {i}: expected a bytes-like object, {} found",
obj.class().slot_name()
))
};
let buffer = PyBuffer::from_object(vm, &obj, BufferFlags::SIMPLE)
.map_err(|_| type_error())?;
if !buffer.desc.is_contiguous() {
return Err(type_error());
}
let item_len = buffer.desc.len;
buffers.push((i, buffer));
item_len
};
objects.push(obj);
len = len.checked_add(item_len).ok_or_else(overflow)?;
if i != 0 {
len = len.checked_add(separator_len).ok_or_else(overflow)?;
}
if len > isize::MAX as usize {
return Err(overflow());
}
if self.len() != count {
return Err(changed());
}
}
let mut joined = Vec::new();
joined
.try_reserve_exact(len)
.map_err(|_| vm.new_memory_error(""))?;
// All Python callbacks have finished. Read the separator only now,
// so same-size mutations made by the iterator or exporters are visible.
let mut buffers_iter = buffers.iter().peekable();
for (i, obj) in objects.iter().enumerate() {
if i != 0 {
// Drop this borrow before acquiring an input buffer: it may
// refer to the same bytearray, with a writer waiting on it.
joined.extend_from_slice(&separator());
}
if buffers_iter.peek().is_some_and(|(index, _)| *index == i) {
let (_, buffer) = buffers_iter.next().unwrap();
joined.extend_from_slice(&buffer.as_contiguous().unwrap());
} else {
joined.extend_from_slice(obj.downcast_ref::<PyBytes>().unwrap().as_bytes());
}
}
Ok(joined)
}
}

/// "y*": any bytes-like object, and nothing else.
impl<'a> TryFromBorrowedObject<'a> for PyBytesInner {
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &'a PyObject) -> PyResult<Self> {
Expand Down Expand Up @@ -624,25 +740,6 @@ impl PyBytesInner {
.py_count(needle.as_slice(), range, |h, n| h.find_iter(n).count()))
}

// stringlib_bytes_join
pub fn join(&self, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
// `PySequence_Fast()`, as in `PyUnicode_Join()`.
let iterable = ArgIterable::<PyObjectRef>::try_from_object(vm, iterable)
.map_err(|_| vm.new_type_error("can only join an iterable"))?;
let iter = iterable.iter_sized(vm)?.enumerate().map(|(i, obj)| {
let obj = obj?;
// Whatever the buffer lookup ran into, the item is only ever
// reported as the wrong kind of thing.
Self::try_from_object(vm, obj.clone()).map_err(|_| {
vm.new_type_error(format!(
"sequence item {i}: expected a bytes-like object, {} found",
obj.class().slot_name()
))
})
});
self.elements.py_join(iter)
}

#[inline]
pub fn find<F>(
&self,
Expand Down
23 changes: 23 additions & 0 deletions extra_tests/snippets/builtin_bytearray.py
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,29 @@ class B(bytearray):
b"123abcxyz"
)

join_separator = bytearray(b",")


def join_mutate_separator():
yield b"a"
join_separator[0] = ord(":")
yield b"b"


assert join_separator.join(join_mutate_separator()) == b"a:b"


def join_resize_separator():
yield b"a"
join_separator.extend(b"!")
yield b"b"


with assert_raises(BufferError):
join_separator.join(join_resize_separator())
join_separator.extend(b"!")
assert join_separator == b":!"


# endswith startswith
assert bytearray(b"abcde").endswith(b"de")
Expand Down
71 changes: 71 additions & 0 deletions extra_tests/snippets/builtin_bytes.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,77 @@
assert b"abc".join((b"123", b"xyz")) == b"123abcxyz"


class JoinBytes(bytes):
def __buffer__(self, flags):
return memoryview(b"override")


class JoinBuffer:
def __init__(self, action=None):
self.action = action

def __buffer__(self, flags):
if self.action is None:
join_chunk[0] = ord("b")
else:
self.action()
return memoryview(b"x")

def __release_buffer__(self, view):
join_events.append("release")


for join_type in (bytes, bytearray):
join_separator = join_type(b",")

def join_mutable_chunks():
chunk = bytearray(b"a")
yield chunk
chunk[:] = b"bc"
yield chunk

assert join_separator.join(join_mutable_chunks()) == b"bc,bc"

def join_broken_iterable():
yield 42
yield b"a"
raise RuntimeError("producer failed")

with assert_raises(RuntimeError):
join_separator.join(join_broken_iterable())

for sequence_type in (list, tuple):

class JoinSequence(sequence_type):
def __iter__(self):
return iter([b"a", b"b"])

assert join_separator.join(JoinSequence([42])) == b"a,b"

result = join_separator.join([JoinBytes(b"original")])
assert result == b"override" and type(result) is join_type

join_chunk = bytearray(b"a")
join_events = []
assert join_separator.join([join_chunk, JoinBuffer()]) == b"b,x"
assert join_events == ["release"]
with assert_raises(TypeError):
join_separator.join([join_chunk, JoinBuffer(lambda: join_chunk.extend(b"c"))])
join_chunk.extend(b"c")
join_items = []
join_exporter = JoinBuffer(join_items.clear)
join_items[:] = [join_exporter, b"y"]
with assert_raises(RuntimeError):
join_separator.join(join_items)
join_events.clear()
with assert_raises(TypeError):
join_separator.join([JoinBuffer(), 42])
assert join_events == ["release"]

join_single = b"single item"
assert b",".join([join_single]) is join_single


# endswith startswith
assert b"abcde".endswith(b"de")
assert b"abcde".endswith(b"")
Expand Down
11 changes: 11 additions & 0 deletions extra_tests/snippets/builtin_str.py
Original file line number Diff line number Diff line change
Expand Up @@ -276,6 +276,17 @@ class JoinStr(str):
assert type("-".join((JoinStr("a"), "b"))) is str
single = "single"
assert "-".join([single]) is single


def join_broken_iterable():
yield 42
yield "a"
raise RuntimeError("producer failed")


with assert_raises(RuntimeError):
"-".join(join_broken_iterable())

assert "HALLO".isupper()
assert not "123".isupper()
assert not "123".islower()
Expand Down
Loading