66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
|
class BitStream:
|
||
|
|
def __init__(self, data: array = []):
|
||
|
|
self.bits_used = 0
|
||
|
|
self.read_offset = 0
|
||
|
|
self.data = data
|
||
|
|
|
||
|
|
def write(self, c_bytes): # ?bytes_to_write
|
||
|
|
bytes_to_write = len(c_bytes)
|
||
|
|
if not len(c_bytes):
|
||
|
|
return
|
||
|
|
|
||
|
|
self.write_bits(c_bytes, bytes_to_write << 3, True)
|
||
|
|
|
||
|
|
def read(self, bytes_to_read: int):
|
||
|
|
bits_to_read = bytes_to_read << 3
|
||
|
|
|
||
|
|
if self.read_offset & 7 == 0:
|
||
|
|
if self.read_offset + bits_to_read > self.bits_used:
|
||
|
|
return False
|
||
|
|
|
||
|
|
read_offset_bytes = self.read_offset >> 3
|
||
|
|
data = self.data[read_offset_bytes:read_offset_bytes + bytes_to_read]
|
||
|
|
|
||
|
|
self.read_offset += bits_to_read
|
||
|
|
return data
|
||
|
|
return self.read_bits(bits_to_read)
|
||
|
|
|
||
|
|
def read_bits(self, bits_to_read: int, align_right: bool = False):
|
||
|
|
if bits_to_read <= 0:
|
||
|
|
return False
|
||
|
|
|
||
|
|
if self.read_offset + bits_to_read > self.bits_used:
|
||
|
|
return False
|
||
|
|
|
||
|
|
offset = 0
|
||
|
|
bits_read_offset_mod8 = self.read_offset & 7
|
||
|
|
|
||
|
|
|
||
|
|
def write_bits(self, c_bytes, bits_to_write: int, align_right: bool = False):
|
||
|
|
if not bits_to_write:
|
||
|
|
return
|
||
|
|
|
||
|
|
offset = 0
|
||
|
|
bits_used_mod8 = self.bits_used & 7
|
||
|
|
|
||
|
|
while bits_to_write:
|
||
|
|
byte = c_bytes[offset]
|
||
|
|
|
||
|
|
if bits_to_write < 8 and align_right:
|
||
|
|
byte <<= 8 - bits_to_write
|
||
|
|
|
||
|
|
if bits_used_mod8 == 0:
|
||
|
|
self.data[self.bits_used >> 3] = byte
|
||
|
|
else:
|
||
|
|
self.data[self.bits_used >> 3] |= byte >> bits_used_mod8
|
||
|
|
if 8 - bits_used_mod8 < 8 and 8 - bits_used_mod8 < bits_to_write:
|
||
|
|
self.data[(self.bits_used >> 3) + 1] = byte << (8 - bits_used_mod8)
|
||
|
|
|
||
|
|
if bits_to_write >= 8:
|
||
|
|
self.bits_used += 8
|
||
|
|
else:
|
||
|
|
self.bits_used += bits_to_write
|
||
|
|
|
||
|
|
bits_to_write -= 8
|
||
|
|
|
||
|
|
offset += 1
|