diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index c463489df61..fc046460b70 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -1,6 +1,6 @@ //! Implementation of the python bytearray object. use super::{ - PositionIterInternal, PyBytes, PyDictRef, PyGenericAlias, PyStrRef, PyTuple, PyTupleRef, + PositionIterInternal, PyBytes, PyDictRef, PyGenericAlias, PyInt, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef, iter::builtins_iter, locked_next, }; use crate::{ @@ -12,7 +12,7 @@ use crate::{ bytes_inner::{ ByteInnerFindOptions, ByteInnerHexOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, ByteInnerReplaceOptions, ByteInnerSplitOptions, ByteInnerStripOptions, ByteInnerSub, - ByteInnerTranslateOptions, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode, + ByteInnerTranslateOptions, ByteOr, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode, }, class::PyClassImpl, common::{ @@ -317,6 +317,21 @@ impl PyByteArray { } fn __contains__(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { + // Exact builtins cannot call Python while preparing the needle. The + // read lock alone protects their search; a temporary export would make + // a concurrent resize fail instead of waiting for that lock. + if let Some(bytes) = needle.downcast_ref_if_exact::(vm) { + return Ok(self.borrow_buf().contains_str(bytes.as_bytes())); + } + let int = if needle.class().is(vm.ctx.types.bool_type) { + needle.downcast_ref::() + } else { + needle.downcast_ref_if_exact::(vm) + }; + if let Some(int) = int { + let byte = int.as_bigint().byte_or(vm)?; + return Ok(self.borrow_buf().contains(&byte)); + } // Converting the needle runs Python, which must not resize this bytearray. self.exports.fetch_add(1, Ordering::Release); scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index 2f2793a76cf..1ea800b64db 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/extra_tests/snippets/builtin_bytearray.py @@ -1504,3 +1504,86 @@ def __index__(self): test_bytearray_callback_exports() + + +def test_bytearray_contains_builtin_and_callback_exports(): + for value in (b"", b"a", b"ab\x00\xff", b"\x01"): + target = bytearray(value) + for needle in ( + b"", + b"a", + b"ab", + b"z", + b"\x00", + b"\xff", + 0, + 97, + 255, + False, + True, + ): + assert (needle in target) == (needle in value) + assert target in target + + def must_raise_buffer_error(action): + try: + action() + except BufferError: + return + raise AssertionError("resize was not rejected") + + target = bytearray(b"a") + view = memoryview(target) + try: + assert b"a" in target + assert 97 in target + must_raise_buffer_error(lambda: target.append(0)) + target[:] = b"b" + assert b"b" in target + assert 98 in target + finally: + view.release() + target.append(0) + assert target == b"b\x00" + + class IndexNeedle: + def __index__(self): + must_raise_buffer_error(lambda: target.append(0)) + target[0] = 98 + return 98 + + class BytesNeedle(bytes): + __index__ = IndexNeedle.__index__ + + for needle in (IndexNeedle(), BytesNeedle(b"\xdd")): + target[:] = b"a" + assert needle in target + assert target == b"b" + target.append(0) + assert target == b"b\x00" + + for needle in (-1, 256): + with assert_raises(ValueError): + needle in target + + events = [0, 0] + + class BufferNeedle: + def __buffer__(self, flags): + must_raise_buffer_error(lambda: target.append(0)) + target[0] = 98 + events[0] += 1 + return memoryview(b"b") + + def __release_buffer__(self, view): + must_raise_buffer_error(lambda: target.append(0)) + events[1] += 1 + + target[:] = b"a" + assert BufferNeedle() in target + assert events == [1, 1] + target.append(0) + assert target == b"b\x00" + + +test_bytearray_contains_builtin_and_callback_exports()