| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| import pytest |
| |
| from pyfory.serialization import Buffer |
| from pyfory.tests.core import require_pyarrow |
| from pyfory.tests.test_stream import OneByteStream |
| from pyfory.utils import lazy_import |
| |
| pa = lazy_import("pyarrow") |
| |
| |
| class RecvIntoOnlyStream: |
| def __init__(self, data: bytes): |
| self._data = data |
| self._offset = 0 |
| |
| def recv_into(self, buffer, size=-1): |
| if self._offset >= len(self._data): |
| return 0 |
| view = memoryview(buffer).cast("B") |
| if size < 0 or size > len(view): |
| size = len(view) |
| if size == 0: |
| return 0 |
| read_size = min(1, size, len(self._data) - self._offset) |
| start = self._offset |
| self._offset += read_size |
| view[:read_size] = self._data[start : start + read_size] |
| return read_size |
| |
| |
| class LegacyRecvIntoOnlyStream: |
| def __init__(self, data: bytes): |
| self._data = data |
| self._offset = 0 |
| |
| def recvinto(self, buffer, size=-1): |
| if self._offset >= len(self._data): |
| return 0 |
| view = memoryview(buffer).cast("B") |
| if size < 0 or size > len(view): |
| size = len(view) |
| if size == 0: |
| return 0 |
| read_size = min(1, size, len(self._data) - self._offset) |
| start = self._offset |
| self._offset += read_size |
| view[:read_size] = self._data[start : start + read_size] |
| return read_size |
| |
| |
| class PartialWriteStream: |
| def __init__(self): |
| self._data = bytearray() |
| |
| def write(self, payload): |
| if not payload: |
| return 0 |
| view = memoryview(payload).cast("B") |
| wrote = min(2, len(view)) |
| self._data.extend(view[:wrote]) |
| return wrote |
| |
| def to_bytes(self): |
| return bytes(self._data) |
| |
| |
| def test_buffer(): |
| buffer = Buffer.allocate(8) |
| buffer.write_bool(True) |
| buffer.write_int8(-1) |
| buffer.write_int8(2**7 - 1) |
| buffer.write_int8(-(2**7)) |
| buffer.write_int16(2**15 - 1) |
| buffer.write_int16(-(2**15)) |
| buffer.write_int32(2**31 - 1) |
| buffer.write_int32(-(2**31)) |
| buffer.write_int64(2**63 - 1) |
| buffer.write_int64(-(2**63)) |
| buffer.write_float(1.0) |
| buffer.write_float(-1.0) |
| buffer.write_double(1.0) |
| buffer.write_double(-1.0) |
| buffer.write_bytes(b"") # write empty buffer |
| buffer.write_buffer(b"") # write empty buffer |
| binary = b"b" * 100 |
| buffer.write_bytes(binary) |
| buffer.write_bytes_and_size(binary) |
| print(f"buffer size {buffer.size()}, writer_index {buffer.get_writer_index()}") |
| new_buffer = Buffer(buffer.get_bytes(0, buffer.get_writer_index())) |
| assert new_buffer.read_bool() is True |
| assert new_buffer.read_int8() == -1 |
| assert new_buffer.read_int8() == 2**7 - 1 |
| assert new_buffer.read_int8() == -(2**7) |
| assert new_buffer.read_int16() == 2**15 - 1 |
| assert new_buffer.read_int16() == -(2**15) |
| assert new_buffer.read_int32() == 2**31 - 1 |
| assert new_buffer.read_int32() == -(2**31) |
| assert new_buffer.read_int64() == 2**63 - 1 |
| assert new_buffer.read_int64() == -(2**63) |
| assert new_buffer.read_float() == 1.0 |
| assert new_buffer.read_float() == -1.0 |
| assert new_buffer.read_double() == 1.0 |
| assert new_buffer.read_double() == -1.0 |
| assert new_buffer.read_bytes(0) == b"" |
| assert new_buffer.read_bytes(0) == b"" |
| assert new_buffer.read_bytes(len(binary)) == binary |
| assert new_buffer.read_bytes_and_size() == binary |
| assert new_buffer.hex() == new_buffer.to_pybytes().hex() |
| assert new_buffer[:10].to_pybytes() == new_buffer.to_pybytes()[:10] |
| assert new_buffer[5:30].to_pybytes() == new_buffer.to_pybytes()[5:30] |
| assert new_buffer[-30:].to_pybytes() == new_buffer.to_pybytes()[-30:] |
| for i in range(len(new_buffer)): |
| assert new_buffer[i] == new_buffer.to_pybytes()[i] |
| assert new_buffer[-i + 1] == new_buffer.to_pybytes()[-i + 1] |
| |
| |
| def test_empty_buffer(): |
| writable_buffer = Buffer.allocate(8) |
| for buffer in [ |
| Buffer.allocate(0), |
| Buffer(b""), |
| Buffer.allocate(8).slice(8), |
| Buffer(b"1").slice(1), |
| ]: |
| assert buffer.to_bytes() == b"" |
| assert buffer.to_pybytes() == b"" |
| assert buffer.slice().to_bytes() == b"" |
| assert buffer.hex() == "" |
| writable_buffer.put_int32(0, 10) |
| writable_buffer.put_buffer(0, buffer, 0, 0) |
| writable_buffer.write_buffer(buffer) |
| assert writable_buffer.get_int32(0) == 10 |
| |
| |
| def test_write_varint32(): |
| buf = Buffer.allocate(32) |
| for i in range(1): |
| for j in range(i): |
| buf.write_int8(1) |
| buf.read_int8() |
| check_varuint32(buf, 1, 1) |
| check_varuint32(buf, 1 << 6, 1) |
| check_varuint32(buf, 1 << 7, 2) |
| check_varuint32(buf, 1 << 13, 2) |
| check_varuint32(buf, 1 << 14, 3) |
| check_varuint32(buf, 1 << 20, 3) |
| check_varuint32(buf, 1 << 21, 4) |
| check_varuint32(buf, 1 << 27, 4) |
| check_varuint32(buf, 1 << 28, 5) |
| check_varuint32(buf, 1 << 30, 5) |
| |
| check_varint32(buf, -1) |
| check_varint32(buf, -1 << 6) |
| check_varint32(buf, -1 << 7) |
| check_varint32(buf, -1 << 13) |
| check_varint32(buf, -1 << 14) |
| check_varint32(buf, -1 << 20) |
| check_varint32(buf, -1 << 21) |
| check_varint32(buf, -1 << 27) |
| check_varint32(buf, -1 << 28) |
| check_varint32(buf, -1 << 30) |
| |
| |
| def check_varuint32(buf: Buffer, value: int, bytes_written: int): |
| assert buf.get_writer_index() == buf.get_reader_index() |
| actual_bytes_written = buf.write_var_uint32(value) |
| assert actual_bytes_written == bytes_written |
| varint = buf.read_var_uint32() |
| assert buf.get_writer_index() == buf.get_reader_index() |
| assert value == varint |
| |
| |
| def check_varint32(buf: Buffer, value: int): |
| assert buf.get_writer_index() == buf.get_reader_index() |
| buf.write_varint32(value) |
| varint = buf.read_varint32() |
| assert buf.get_writer_index() == buf.get_reader_index() |
| assert value == varint |
| |
| |
| @require_pyarrow |
| def test_buffer_protocol(): |
| # test buffer protocol compatibility with pyarrow |
| buffer = Buffer.allocate(32) |
| binary = b"b" * 100 |
| buffer.write_bytes_and_size(binary) |
| assert bytes(buffer) == bytes(pa.py_buffer(buffer)) |
| assert buffer.to_bytes() == bytes(pa.py_buffer(buffer)) |
| |
| |
| def test_grow(): |
| binary = b"a" * 10 |
| buffer = Buffer(binary) |
| assert not buffer.own_data() |
| buffer.write_bytes(binary) |
| assert not buffer.own_data() |
| buffer.write_bytes(binary) |
| assert buffer.own_data() |
| |
| |
| def test_write_var_uint64(): |
| buf = Buffer.allocate(32) |
| check_varuint64(buf, -1, 9) |
| for i in range(32): |
| for j in range(i): |
| buf.write_int8(1) |
| buf.read_int8() |
| check_varuint64(buf, -1, 9) |
| check_varuint64(buf, 1, 1) |
| check_varuint64(buf, 1 << 6, 1) |
| check_varuint64(buf, 1 << 7, 2) |
| check_varuint64(buf, -(2**6), 9) |
| check_varuint64(buf, -(2**7), 9) |
| check_varuint64(buf, 1 << 13, 2) |
| check_varuint64(buf, 1 << 14, 3) |
| check_varuint64(buf, -(2**13), 9) |
| check_varuint64(buf, -(2**14), 9) |
| check_varuint64(buf, 1 << 20, 3) |
| check_varuint64(buf, 1 << 21, 4) |
| check_varuint64(buf, -(2**20), 9) |
| check_varuint64(buf, -(2**21), 9) |
| check_varuint64(buf, 1 << 27, 4) |
| check_varuint64(buf, 1 << 28, 5) |
| check_varuint64(buf, -(2**27), 9) |
| check_varuint64(buf, -(2**28), 9) |
| check_varuint64(buf, 1 << 30, 5) |
| check_varuint64(buf, -(2**30), 9) |
| check_varuint64(buf, 1 << 31, 5) |
| check_varuint64(buf, -(2**31), 9) |
| check_varuint64(buf, 1 << 32, 5) |
| check_varuint64(buf, -(2**32), 9) |
| check_varuint64(buf, 1 << 34, 5) |
| check_varuint64(buf, -(2**34), 9) |
| check_varuint64(buf, 1 << 35, 6) |
| check_varuint64(buf, -(2**35), 9) |
| check_varuint64(buf, 1 << 41, 6) |
| check_varuint64(buf, -(2**41), 9) |
| check_varuint64(buf, 1 << 42, 7) |
| check_varuint64(buf, -(2**42), 9) |
| check_varuint64(buf, 1 << 48, 7) |
| check_varuint64(buf, -(2**48), 9) |
| check_varuint64(buf, 1 << 49, 8) |
| check_varuint64(buf, -(2**49), 9) |
| check_varuint64(buf, 1 << 55, 8) |
| check_varuint64(buf, -(2**55), 9) |
| check_varuint64(buf, 1 << 56, 9) |
| check_varuint64(buf, -(2**56), 9) |
| check_varuint64(buf, 1 << 62, 9) |
| check_varuint64(buf, -(2**62), 9) |
| check_varuint64(buf, 1 << 63 - 1, 9) |
| check_varuint64(buf, -(2**63), 9) |
| |
| |
| def check_varuint64(buf: Buffer, value: int, bytes_written: int): |
| assert buf.get_writer_index() == buf.get_reader_index() |
| actual_bytes_written = buf.write_var_uint64(value) |
| assert actual_bytes_written == bytes_written |
| varint = buf.read_var_uint64() |
| assert buf.get_writer_index() == buf.get_reader_index() |
| assert value == varint |
| |
| |
| def test_buffer_flush_stream(): |
| stream = PartialWriteStream() |
| buffer = Buffer.allocate(16) |
| output_stream = Buffer.wrap_output_stream(stream) |
| buffer.bind_output_stream(output_stream) |
| payload = b"stream-flush-buffer" |
| buffer.write_bytes(payload) |
| output_stream.force_flush() |
| assert stream.to_bytes() == payload |
| assert buffer.get_writer_index() == 0 |
| |
| |
| def test_wrap_output_stream_invalid_target_raises(): |
| with pytest.raises(ValueError): |
| Buffer.wrap_output_stream(object()) |
| |
| |
| def test_output_stream_try_flush_preserves_bound_buffer_when_barrier_active(): |
| stream = PartialWriteStream() |
| output_stream = Buffer.wrap_output_stream(stream) |
| buffer = Buffer.allocate(32) |
| buffer.bind_output_stream(output_stream) |
| payload = b"x" * 5000 |
| |
| output_stream.enter_flush_barrier() |
| buffer.write_bytes(payload) |
| output_stream.try_flush() |
| output_stream.try_flush() |
| assert buffer.get_writer_index() == len(payload) |
| assert stream.to_bytes() == b"" |
| |
| output_stream.exit_flush_barrier() |
| output_stream.try_flush() |
| assert buffer.get_writer_index() == 0 |
| |
| output_stream.force_flush() |
| assert stream.to_bytes() == payload |
| |
| |
| def test_output_stream_try_flush_small_payload_needs_force_flush(): |
| stream = PartialWriteStream() |
| output_stream = Buffer.wrap_output_stream(stream) |
| buffer = Buffer.allocate(32) |
| buffer.bind_output_stream(output_stream) |
| payload = b"small-payload" |
| buffer.write_bytes(payload) |
| |
| output_stream.try_flush() |
| assert buffer.get_writer_index() == len(payload) |
| assert stream.to_bytes() == b"" |
| |
| output_stream.force_flush() |
| assert buffer.get_writer_index() == 0 |
| assert stream.to_bytes() == payload |
| |
| |
| def test_write_buffer(): |
| buf = Buffer.allocate(32) |
| buf.write(b"") |
| buf.write(b"123") |
| buf.write(Buffer.allocate(32)) |
| assert buf.get_writer_index() == 35 |
| assert buf.read(0) == b"" |
| assert buf.read(3) == b"123" |
| |
| |
| def test_read_bytes_as_int64(): |
| # test small buffer whose length < 8 |
| buf = Buffer(b"1234") |
| assert buf.read_bytes_as_int64(0) == 0 |
| assert buf.read_bytes_as_int64(1) == 49 |
| |
| # test big buffer whose length > 8 |
| buf = Buffer(b"12345678901234") |
| assert buf.read_bytes_as_int64(0) == 0 |
| assert buf.read_bytes_as_int64(1) == 49 |
| assert buf.read_bytes_as_int64(8) == 4123106164818064178 |
| |
| # test fix for `OverflowError: Python int too large to convert to C long` |
| buf = Buffer(b"\xa6IOr\x9ch)\x80\x12\x02") |
| buf.read_bytes_as_int64(8) |
| |
| |
| def test_stream_buffer_read(): |
| writer = Buffer.allocate(32) |
| writer.write_uint32(0x01020304) |
| writer.write_int64(-1234567890) |
| writer.write_var_uint32(300) |
| writer.write_varint64(-4567890123) |
| writer.write_tagged_uint64(0x123456789) |
| writer.write_var_uint64(0x1FFFF) |
| writer.write_bytes_and_size(b"stream-data") |
| writer.write_string("hello-stream") |
| |
| data = writer.get_bytes(0, writer.get_writer_index()) |
| stream = OneByteStream(data) |
| reader = Buffer.from_stream(stream) |
| |
| assert reader.read_uint32() == 0x01020304 |
| assert reader.read_int64() == -1234567890 |
| assert reader.read_var_uint32() == 300 |
| assert reader.read_varint64() == -4567890123 |
| assert reader.read_tagged_uint64() == 0x123456789 |
| assert reader.read_var_uint64() == 0x1FFFF |
| assert reader.read_bytes_and_size() == b"stream-data" |
| assert reader.read_string() == "hello-stream" |
| |
| |
| def test_stream_buffer_read_with_recv_into(): |
| reader = Buffer.from_stream(RecvIntoOnlyStream(bytes([0x11, 0x22, 0x33, 0x44]))) |
| assert reader.read_uint32() == 0x44332211 |
| |
| |
| def test_stream_buffer_read_with_legacy_recvinto(): |
| reader = Buffer.from_stream(LegacyRecvIntoOnlyStream(bytes([0x11, 0x22, 0x33, 0x44]))) |
| assert reader.read_uint32() == 0x44332211 |
| |
| |
| def test_stream_buffer_set_reader_index(): |
| reader = Buffer.from_stream(OneByteStream(bytes([0x11, 0x22, 0x33, 0x44, 0x55]))) |
| reader.set_reader_index(4) |
| assert reader.read_uint8() == 0x55 |
| |
| |
| def test_stream_buffer_set_reader_index_out_of_bound(): |
| reader = Buffer.from_stream(OneByteStream(b"\x11\x22\x33")) |
| with pytest.raises(Exception, match="Buffer out of bound"): |
| reader.set_reader_index(10) |
| |
| |
| def test_stream_buffer_read_bytes_and_skip_update_reader_index(): |
| reader = Buffer.from_stream(OneByteStream(bytes(range(20))), buffer_size=2) |
| assert reader.read_bytes(5) == bytes([0, 1, 2, 3, 4]) |
| assert reader.get_reader_index() == 5 |
| reader.skip(5) |
| assert reader.get_reader_index() == 10 |
| |
| |
| def test_stream_buffer_short_read_error(): |
| reader = Buffer.from_stream(OneByteStream(b"\x01\x02\x03")) |
| with pytest.raises(Exception, match="Buffer out of bound"): |
| reader.read_uint32() |
| |
| |
| if __name__ == "__main__": |
| test_grow() |
