byte-buffered python BitStream

bitstream.py

from io import BytesIO


class BitStream:
    def __init__(self, data: BytesIO):
        self.data = data
        self.buffer = 0
        self.index = 0
        self.data.seek(0)

    def read_bits(self, bits: int) -> int:
        bits_to_read = bits
        result: int = 0

        if 0 != self.index % 8:
            buffer_left = -self.index % 8
            bits_to_read -= buffer_left
            result |= self.buffer & ((1 << buffer_left) - 1)

        self.index += bits
        bytes_to_read = (bits_to_read + 7) // 8
        got_bytes = self.data.read(bytes_to_read)
        if got_bytes == b'' and bytes_to_read != 0:
            # reached EOF
            return -1

        full_bytes = got_bytes[:bits_to_read // 8]

        result <<= bits_to_read // 8
        result |= int.from_bytes(full_bytes, byteorder='big', signed=False)

        if 0 != bits_to_read % 8:
            bits_to_read %= 8
            last_byte = got_bytes[bytes_to_read - 1]
            self.buffer = last_byte
            result <<= bits_to_read
            result |= (last_byte >> (-bits_to_read % 8)) & ((1 << bits_to_read) - 1)

        return result


if __name__ == "__main__":
    x = BitStream(BytesIO(bytes([0b11000011, 0b10101010, 0b11101110, 0b01100110])))
    assert bin(r := x.read_bits(4)) == "0b1100", bin(r)
    assert x.index == 4, x.index
    assert x.buffer != 0x00

    assert bin(r := x.read_bits(4)) == "0b11", bin(r)
    assert x.index == 8, x.index
    # assert x.buffer == 0x00

    assert bin(r := x.read_bits(5)) == "0b10101", bin(r)
    assert x.index == 8 + 5, x.index
    assert x.buffer != 0x00

    assert bin(r := x.read_bits(3)) == "0b10", bin(r)
    assert x.index == 8 + 8, x.index

    assert bin(r := x.read_bits(12)) == "0b111011100110", bin(r)
    assert x.index == 8 + 8 + 12, x.index

    assert bin(r := x.read_bits(4)) == "0b110", bin(r)
    assert x.index == 8 + 8 + 12 + 4, x.index

    # EOF test
    assert bin(r := x.read_bits(1)) == "-0b1", bin(r)
    assert x.index == 8 + 8 + 12 + 4 + 1, x.index

    test = 0b1011_01010001_01010010_10100100_11010100_01010101_10101001
    test_bytes = int.to_bytes(test, (test.bit_length() + 7) // 8, 'big', signed=False)
    print(test_bytes)

    x = BitStream(BytesIO(test_bytes))
    for i in range((test.bit_length() + 7) // 8):
        print(f"{x.read_bits(8):08b}")