peterlu14
190 lines
4.2 KiB
Python
190 lines
4.2 KiB
Python
import abc
|
|
import struct
|
|
from typing import List, IO, Optional
|
|
|
|
from biopro.util.address import ADDRESS
|
|
|
|
|
|
class AbstractEncoder(metaclass=abc.ABCMeta):
|
|
__slots__ = ()
|
|
|
|
@abc.abstractmethod
|
|
def write(self, data: bytes):
|
|
pass
|
|
|
|
def u8(self, v: int):
|
|
self.write(struct.pack('B', v))
|
|
|
|
def u16(self, v: int):
|
|
self.write(struct.pack('>H', v))
|
|
|
|
def u32(self, v: int):
|
|
self.write(struct.pack('>I', v))
|
|
|
|
def u64(self, v: int):
|
|
self.write(struct.pack('>Q', v))
|
|
|
|
def u8a(self, v: List[int]):
|
|
self.write(struct.pack('%dB' % len(v), *v))
|
|
|
|
def u16a(self, v: List[int]):
|
|
self.write(struct.pack('%dH' % len(v), *v))
|
|
|
|
def i8(self, v: int):
|
|
self.write(struct.pack('b', v))
|
|
|
|
def i16(self, v: int):
|
|
self.write(struct.pack('>h', v))
|
|
|
|
def i32(self, v: int):
|
|
self.write(struct.pack('>i', v))
|
|
|
|
def i64(self, v: int):
|
|
self.write(struct.pack('>q', v))
|
|
|
|
def i8a(self, v: List[int]):
|
|
self.write(struct.pack('%db' % len(v), *v))
|
|
|
|
def encode_string(self, s: Optional[str]):
|
|
if s is None:
|
|
self.write(struct.pack('B', 0))
|
|
else:
|
|
d = s.encode()
|
|
self.write(struct.pack('B%ds' % len(d), len(d), d))
|
|
|
|
def encode_address(self, a: ADDRESS):
|
|
self.write(struct.pack('6B', *reversed(a)))
|
|
|
|
|
|
class AbstractDecoder(metaclass=abc.ABCMeta):
|
|
__slots__ = ()
|
|
|
|
@abc.abstractmethod
|
|
def read(self, size: int) -> bytes:
|
|
pass
|
|
|
|
def u8(self) -> int:
|
|
return ord(self.read(1))
|
|
|
|
def u16(self) -> int:
|
|
return struct.unpack('>H', self.read(2))[0]
|
|
|
|
def u32(self) -> int:
|
|
return struct.unpack('>I', self.read(4))[0]
|
|
|
|
def u64(self) -> int:
|
|
return struct.unpack('>Q', self.read(8))[0]
|
|
|
|
def u8a(self, size: int) -> List[int]:
|
|
return [ord(self.read(1)) for _ in range(size)]
|
|
|
|
def i8(self) -> int:
|
|
return struct.unpack('b', self.read(1))[0]
|
|
|
|
def i16(self) -> int:
|
|
return struct.unpack('>h', self.read(2))[0]
|
|
|
|
def i32(self) -> int:
|
|
return struct.unpack('>i', self.read(4))[0]
|
|
|
|
def i64(self) -> int:
|
|
return struct.unpack('>q', self.read(8))[0]
|
|
|
|
def i8a(self, size: int) -> List[int]:
|
|
return list(struct.unpack('%db' % size, self.read(size)))
|
|
|
|
def decode_string(self) -> str:
|
|
s = ord(self.read(1))
|
|
|
|
if s == 0:
|
|
return ''
|
|
else:
|
|
return self.read(s).decode()
|
|
|
|
def decode_string_nullable(self) -> Optional[str]:
|
|
s = ord(self.read(1))
|
|
|
|
if s == 0:
|
|
return None
|
|
else:
|
|
return self.read(s).decode()
|
|
|
|
def decode_address(self) -> ADDRESS:
|
|
# noinspection PyTypeChecker
|
|
return tuple(reversed(struct.unpack('6B', self.read(6))))
|
|
|
|
|
|
class FileEncoder(AbstractEncoder):
|
|
__slots__ = ('_io',)
|
|
|
|
def __init__(self, io: IO):
|
|
if not io.writable():
|
|
raise RuntimeError('file not writable')
|
|
|
|
self._io = io
|
|
|
|
def write(self, data: bytes):
|
|
self._io.write(data)
|
|
|
|
def flush(self):
|
|
self._io.flush()
|
|
|
|
|
|
class FileDecoder(AbstractDecoder):
|
|
__slots__ = ('_io',)
|
|
|
|
def __init__(self, io: IO):
|
|
if not io.readable():
|
|
raise RuntimeError('file not readable')
|
|
|
|
self._io = io
|
|
|
|
def read(self, size: int) -> bytes:
|
|
ret = self._io.read(size)
|
|
|
|
if len(ret) != size:
|
|
raise EOFError
|
|
|
|
return ret
|
|
|
|
|
|
class BytesEncoder(AbstractEncoder):
|
|
__slots__ = ('_buffer',)
|
|
|
|
def __init__(self):
|
|
self._buffer = []
|
|
|
|
def write(self, data: bytes):
|
|
self._buffer.append(data)
|
|
|
|
def get(self) -> bytes:
|
|
return b''.join(self._buffer)
|
|
|
|
|
|
class BytesDecoder(AbstractDecoder):
|
|
__slots__ = ('_offset', '_data')
|
|
|
|
def __init__(self, data: bytes):
|
|
self._offset = 0
|
|
self._data = data
|
|
|
|
def read(self, size: int) -> bytes:
|
|
off = self._offset
|
|
ret = self._data[off:off + size]
|
|
|
|
if len(ret) != size:
|
|
raise EOFError
|
|
|
|
self._offset = off + size
|
|
|
|
return ret
|
|
|
|
def reset(self, data: bytes, append=False):
|
|
if append:
|
|
self._data = self._data[self._offset:] + data
|
|
|
|
else:
|
|
self._data = data
|
|
|
|
self._offset = 0
|