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
48 changes: 47 additions & 1 deletion crates/vm/src/builtins/str.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use super::{
PositionIterInternal, PyBytesRef, PyDict, PyTupleRef, PyType, PyTypeRef,
PositionIterInternal, PyBytesRef, PyDict, PyList, PyTuple, PyTupleRef, PyType, PyTypeRef,
int::{PyInt, PyIntRef},
iter::{IterStatus, builtins_iter},
};
Expand Down Expand Up @@ -1217,6 +1217,13 @@ impl PyStr {

#[pymethod]
fn join(zelf: PyRef<Self>, iterable: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyStrRef> {
// `PySequence_Fast()` hands a list or tuple over as-is.
if let Some(list) = iterable.downcast_ref_if_exact::<PyList>(vm) {
return Self::join_items(&zelf, &list.borrow_vec(), vm);
}
if let Some(tuple) = iterable.downcast_ref_if_exact::<PyTuple>(vm) {
return Self::join_items(&zelf, tuple.as_slice(), vm);
}
// `PyUnicode_Join()` reaches its elements through `PySequence_Fast()`,
// which fills a list from the iterator and so asks it how long it is,
// and which has its own wording for what it cannot iterate.
Expand All @@ -1243,6 +1250,45 @@ impl PyStr {
Ok(vm.ctx.new_str(joined))
}

/// `join` over already-materialized items: checks them and sizes the result before copying.
fn join_items(
zelf: &Py<Self>,
items: &[PyObjectRef],
vm: &VirtualMachine,
) -> PyResult<PyStrRef> {
fn item_str<'a>(
i: usize,
obj: &'a PyObject,
vm: &VirtualMachine,
) -> PyResult<&'a Py<PyStr>> {
obj.downcast_ref::<PyStr>().ok_or_else(|| {
vm.new_type_error(format!(
"sequence item {i}: expected str instance, {} found",
obj.class().slot_name()
))
})
}
let sep = zelf.as_wtf8();
let mut len = sep.len().saturating_mul(items.len().saturating_sub(1));
for (i, obj) in items.iter().enumerate() {
len = len.saturating_add(item_str(i, obj, vm)?.as_wtf8().len());
}
if let [only] = items {
let only = item_str(0, only, vm)?;
if only.class().is(vm.ctx.types.str_type) {
return Ok(only.to_owned());
}
}
let mut joined = Wtf8Buf::with_capacity(len);
for (i, obj) in items.iter().enumerate() {
if i > 0 {
joined.push_wtf8(sep);
}
joined.push_wtf8(item_str(i, obj, vm)?.as_wtf8());
}
Ok(vm.ctx.new_str(joined))
}

/// The bytes the character range `range` spans and the byte offset it
/// starts at, or `None` if the range is inverted.
///
Expand Down
16 changes: 16 additions & 0 deletions extra_tests/snippets/builtin_str.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,22 @@
)
assert "abc\t12345\txyz".expandtabs() == "abc 12345 xyz"
assert "-".join(["1", "2", "3"]) == "1-2-3"
assert "-".join(("1", "2")) == "1-2"
assert "-".join([]) == ""
assert "-".join(x for x in "ab") == "a-b"
with assert_raises(TypeError) as cm:
"-".join(["a", 1])
assert str(cm.exception) == "sequence item 1: expected str instance, int found"


class JoinStr(str):
pass


assert type("-".join([JoinStr("a")])) is str
assert type("-".join((JoinStr("a"), "b"))) is str
single = "single"
assert "-".join([single]) is single
assert "HALLO".isupper()
assert not "123".isupper()
assert not "123".islower()
Expand Down
Loading