class BitWriter: def __init__(self) -> None: self._bits: list[int] = [] @property def bit_length(self) -> int: return len(self._bits) def write(self, value: int, width: int) -> None: if width < 0: raise ValueError("width must be positive") if value < 0 or value >= (1 << width): raise ValueError(f"value {value} does not fit in {width} bits") for shift in range(width - 1, -1, -1): self._bits.append((value >> shift) & 1) def to_bytes(self) -> bytes: padding = (-len(self._bits)) % 8 bits = self._bits + [0] * padding out = bytearray() for i in range(0, len(bits), 8): byte = 0 for bit in bits[i : i + 8]: byte = (byte << 1) | bit out.append(byte) return bytes(out) class BitReader: def __init__(self, data: bytes, bit_length: int) -> None: self._bits: list[int] = [] for byte in data: for shift in range(7, -1, -1): self._bits.append((byte >> shift) & 1) self._bits = self._bits[:bit_length] self._pos = 0 def read(self, width: int) -> int: if self._pos + width > len(self._bits): raise ValueError("payload ended unexpectedly") value = 0 for bit in self._bits[self._pos : self._pos + width]: value = (value << 1) | bit self._pos += width return value @property def remaining(self) -> int: return len(self._bits) - self._pos