"""Read-write utility functions.""" import struct import typing class EndianIO: # TODO(Sybren): note as UCHAR: struct.Struct = None and move actual structs to LittleEndianTypes UCHAR = struct.Struct(b' int: """Write a (truncated) string as UTF-8. The string will always be written 0-terminated. :param fileobj: the file to write to. :param astring: the string to write. :param fieldlen: the field length in bytes. :returns: the number of bytes written. """ assert isinstance(astring, str) # TODO: truncate the string on a UTF-8 character boundary to avoid creating invalid UTF-8. encoded = astring.encode('utf-8')[:fieldlen-1] + b'\0' return fileobj.write(encoded) @classmethod def write_bytes(cls, fileobj: typing.BinaryIO, data: bytes, fieldlen: int) -> int: """Write (truncated) bytes. When len(data) < fieldlen, a terminating b'\0' will be appended. :returns: the number of bytes written. """ assert isinstance(data, (bytes, bytearray)) if len(data) >= fieldlen: to_write = data[0:fieldlen] else: to_write = data + b'\0' return fileobj.write(to_write) @classmethod def read_bytes0(cls, fileobj, length): data = fileobj.read(length) return cls.read_data0(data) @classmethod def read_data0_offset(cls, data, offset): add = data.find(b'\0', offset) - offset return data[offset:offset + add] @classmethod def read_data0(cls, data): add = data.find(b'\0') if add < 0: return data return data[:add] class LittleEndianTypes(EndianIO): pass class BigEndianTypes(LittleEndianTypes): UCHAR = struct.Struct(b'>B') USHORT = struct.Struct(b'>H') USHORT2 = struct.Struct(b'>HH') # two shorts in a row SSHORT = struct.Struct(b'>h') UINT = struct.Struct(b'>I') SINT = struct.Struct(b'>i') FLOAT = struct.Struct(b'>f') ULONG = struct.Struct(b'>Q')