diff --git a/tests/snippets/bytearray.py b/tests/snippets/bytearray.py index 4cf8e608c..889c34056 100644 --- a/tests/snippets/bytearray.py +++ b/tests/snippets/bytearray.py @@ -271,6 +271,9 @@ assert bytearray(b"").join( ) == bytearray(b"jiljlkmoomkaaaa") with assert_raises(TypeError): bytearray(b"").join((b"km", "kl")) +assert bytearray(b"abc").join(( + bytearray(b"123"), bytearray(b"xyz") +)) == bytearray(b"123abcxyz") # endswith startswith diff --git a/tests/snippets/bytes.py b/tests/snippets/bytes.py index 33774fe2c..69a8f8e9f 100644 --- a/tests/snippets/bytes.py +++ b/tests/snippets/bytes.py @@ -270,6 +270,8 @@ assert ( with assert_raises(TypeError): b"".join((b"km", "kl")) +assert b"abc".join((b"123", b"xyz")) == b"123abcxyz" + # endswith startswith assert b"abcde".endswith(b"de") diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs index 0597863e7..9e4936c88 100644 --- a/vm/src/obj/objbytearray.rs +++ b/vm/src/obj/objbytearray.rs @@ -290,7 +290,7 @@ impl PyByteArrayRef { } #[pymethod(name = "join")] - fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { + fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { self.inner.borrow().join(iter, vm) } diff --git a/vm/src/obj/objbyteinner.rs b/vm/src/obj/objbyteinner.rs index 09db11ed5..983e8a4ed 100644 --- a/vm/src/obj/objbyteinner.rs +++ b/vm/src/obj/objbyteinner.rs @@ -826,11 +826,14 @@ impl PyByteInner { Ok(total) } - pub fn join(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { + pub fn join(&self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { let mut refs = vec![]; for v in iter.iter(vm)? { let v = v?; - refs.extend(PyByteInner::try_from_object(vm, v)?.elements) + if !refs.is_empty() { + refs.extend(&self.elements); + } + refs.extend(v.elements); } Ok(vm.ctx.new_bytes(refs)) diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs index 61a1595f7..2e2a0d6db 100644 --- a/vm/src/obj/objbytes.rs +++ b/vm/src/obj/objbytes.rs @@ -263,7 +263,7 @@ impl PyBytesRef { } #[pymethod(name = "join")] - fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { + fn join(self, iter: PyIterable, vm: &VirtualMachine) -> PyResult { self.inner.join(iter, vm) }