timeit benchmark for different implementation of the pipe-like stream reader writer

test_stream_reader_writer.py

from io import BytesIO


class BytesIOReaderWriter:
    _buf: BytesIO

    def __init__(self):
        self._buf = BytesIO()

    def write(self, data):
        self._buf.write(data)

    def read(self, n=-1):
        # always adjust the position to the beginning
        self._buf.seek(0)
        data = self._buf.read(n)
        # copy the remaining data to a new buf
        # doing this because no decent way to resize the buffer
        orig_data = self._buf.getvalue()
        self._buf = BytesIO(orig_data[len(data):])
        return data


class BytesReaderWriter:
    _bytes: bytes

    def __init__(self):
        self._bytes = b''

    def write(self, data):
        self._bytes = self._bytes + data

    def read(self, n):
        data = self._bytes[:n]
        self._bytes = self._bytes[n:]
        return data


class ListReaderWriter:
    _queue: list

    def __init__(self):
        self._queue = []

    def write(self, data):
        data_bytes_list = [data[i:i+1] for i in range(len(data))]
        self._queue += data_bytes_list

    def read(self, n):
        data = b''.join(self._queue[:n])
        self._queue = self._queue[n:]
        return data


def benchmark(rwtor):
    # rwtor = BytesIO()
    data = b"abcd"
    rwtor.write(data[:1])
    rwtor.write(data[1:3])
    assert rwtor.read(2) == data[:2]
    assert rwtor.read(1) == data[2:3]
    rwtor.write(data[3:])
    assert rwtor.read(1) == data[3:]


def test_bytesio_rwtor():
    benchmark(BytesIOReaderWriter())


def test_bytes_rwtor():
    benchmark(BytesReaderWriter())


def test_list_rwtor():
    benchmark(ListReaderWriter())


if __name__ == "__main__":
    import timeit
    test_funcs = (
        test_bytesio_rwtor,
        test_bytes_rwtor,
        test_list_rwtor,
    )
    for func in test_funcs:
        func_name = func.__name__
        result = timeit.timeit(
            f"{func_name}()",
            setup=f"from __main__ import {func_name}",
        )
        print(f"{func_name}: {result}")