BitStream.py

class BitStream:
    def __init__(self, data: bytes):
        self.data = data
        self.bitOffset = 0

    def read_bits(self, bits: int):
        if (bits + self.bitOffset) // 8 > len(self.data):
            raise EOFError()

        res = 0
        for i in range(bits):
            byte = self.data[self.bitOffset // 8]
            bit = (byte >> self.bitOffset % 8) & 1
            res |= bit << i
            self.bitOffset += 1
        return res