68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
from array import array
|
|
|
|
class Bitstream:
|
|
def __init__(self):
|
|
self.offset = 0
|
|
self.buffer = array("B")
|
|
|
|
# This is not finished.. would like to implement bit_length support
|
|
def write(self, value: bool):
|
|
offset = self.offset % 8
|
|
|
|
if self.offset // 8 == len(self.buffer):
|
|
self.buffer.append(0)
|
|
|
|
self.offset += 1
|
|
|
|
if value == False:
|
|
return
|
|
|
|
mask = 1 << (7 - offset)
|
|
self.buffer[-1] |= mask
|
|
|
|
"""
|
|
def write(self, value: int, bit_length: int):
|
|
offset = self.offset % 8
|
|
|
|
shifts = offset + bit_length
|
|
while shifts > 8:
|
|
|
|
|
|
if offset == 0:
|
|
self.buffer.append(0)
|
|
|
|
v = self.buffer[buffer_index]
|
|
|
|
mask = (1 << ((self.offset % 8) + 1))
|
|
"""
|
|
|
|
|
|
def read(self, bit_length: int) -> int:
|
|
byte_read_from = self.offset // 8
|
|
byte_read_to = (self.offset + bit_length - 1) // 8
|
|
bytes_to_read = byte_read_to - byte_read_from + 1
|
|
|
|
if byte_read_to >= len(self.buffer):
|
|
raise Exception("Reading beyond the buffer")
|
|
|
|
value = 0
|
|
shifts = bytes_to_read - 1
|
|
while shifts >= 0:
|
|
value += self.buffer[byte_read_to - shifts] << (shifts * 8)
|
|
shifts -= 1
|
|
|
|
value <<= self.offset % 8
|
|
value &= (1 << (8 * bytes_to_read)) - 1
|
|
value >>= (bytes_to_read * 8) - bit_length
|
|
|
|
self.offset += bit_length
|
|
return value
|
|
|
|
def __repr__(self) -> str:
|
|
b = bytearray(" " + " ".join(" ".join(format(c, "08b")) for c in self.buffer) + " ", "ascii")
|
|
m = self.offset * 2 + (self.offset // 8) - (self.offset // (max(1, len(self.buffer)) * 8))
|
|
#b[m] = 91 # [
|
|
#b[m + 2] = 93 # ]
|
|
b[m] = 124 # |
|
|
return b.decode()
|