Some type hint fixes for protocol necessary to compile with mypyc.
This commit is contained in:
Jennifer Taylor
parent
48405f00c5
commit
4a1579b099
@ -126,7 +126,7 @@ class DXTBuffer:
|
||||
pixel_color = (0, 0, 0, alpha)
|
||||
|
||||
# While not surpassing the image dimensions, assign pixels the colors.
|
||||
if (x + i) < self.width and (y + j) < self.height:
|
||||
if pixel_color is not None and (x + i) < self.width and (y + j) < self.height:
|
||||
self.decompressed_buffer[(y + j) * self.width + (x + i)] = (
|
||||
struct.pack('<BBBB', *pixel_color)
|
||||
)
|
||||
|
||||
@ -51,7 +51,7 @@ class PackedOrdering:
|
||||
size - Number of bytes to work with as an integer
|
||||
allow_expansion - Boolean describing whether to add to the end of the order when needed
|
||||
"""
|
||||
self.order: List[int] = []
|
||||
self.order: List[Optional[int]] = []
|
||||
self.expand = allow_expansion
|
||||
|
||||
for i in range(size):
|
||||
@ -269,6 +269,8 @@ class BinaryDecoder:
|
||||
A string representing the name in ascii
|
||||
"""
|
||||
length = self.stream.read_int()
|
||||
if length is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read node name length!")
|
||||
binary_length = int(((length * 6) + 7) / 8)
|
||||
|
||||
def int_to_bin(integer: int) -> str:
|
||||
@ -280,7 +282,10 @@ class BinaryDecoder:
|
||||
|
||||
data = ''
|
||||
for i in range(binary_length):
|
||||
data = data + int_to_bin(self.stream.read_int())
|
||||
next_byte = self.stream.read_int()
|
||||
if next_byte is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read node name!")
|
||||
data = data + int_to_bin(next_byte)
|
||||
data_str = [data[i:(i + 6)] for i in range(0, len(data), 6)]
|
||||
data_int = [int(val, 2) for val in data_str]
|
||||
ret = ''.join([Node.NODE_NAME_CHARS[val] for val in data_int])
|
||||
@ -301,6 +306,8 @@ class BinaryDecoder:
|
||||
|
||||
while True:
|
||||
child_type = self.stream.read_int()
|
||||
if child_type is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read node type!")
|
||||
|
||||
if child_type == Node.END_OF_NODE:
|
||||
return node
|
||||
@ -320,13 +327,18 @@ class BinaryDecoder:
|
||||
Node object
|
||||
"""
|
||||
if self.executed:
|
||||
raise Exception("Logic error, should only call this once per instance")
|
||||
raise BinaryEncodingException("Logic error, should only call this once per instance")
|
||||
self.executed = True
|
||||
|
||||
# Read the header first
|
||||
header_length = self.stream.read_int(4)
|
||||
if header_length is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read header length!")
|
||||
|
||||
root = self.__read_node(self.stream.read_int())
|
||||
node_type = self.stream.read_int()
|
||||
if node_type is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read root node type!")
|
||||
root = self.__read_node(node_type)
|
||||
|
||||
eod = self.stream.read_int()
|
||||
if eod != Node.END_OF_DOCUMENT:
|
||||
@ -378,6 +390,8 @@ class BinaryDecoder:
|
||||
loc = ordering.get_next_short()
|
||||
elif alignment == 4:
|
||||
loc = ordering.get_next_int()
|
||||
if loc is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read node data location!")
|
||||
|
||||
if size is None:
|
||||
# The size should be read from the first 4 bytes
|
||||
@ -419,6 +433,8 @@ class BinaryDecoder:
|
||||
else:
|
||||
# Array value
|
||||
loc = ordering.get_next_int()
|
||||
if loc is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to read array length location!")
|
||||
|
||||
# The raw size in bytes
|
||||
length = struct.unpack('>I', body[loc:(loc + 4)])[0]
|
||||
@ -602,6 +618,8 @@ class BinaryEncoder:
|
||||
loc = ordering.get_next_short()
|
||||
elif alignment == 4:
|
||||
loc = ordering.get_next_int()
|
||||
if loc is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting to allocate node location!")
|
||||
|
||||
if dtype == 'str':
|
||||
# Need to convert this to encoding from standard string.
|
||||
@ -633,6 +651,9 @@ class BinaryEncoder:
|
||||
continue
|
||||
elif composite:
|
||||
# Array, but not, somewhat silly
|
||||
if size is None:
|
||||
raise Exception("Logic error, node size not set yet this is not an attribute!")
|
||||
|
||||
encode_value = f'>{enc}'
|
||||
self.__add_data(struct.pack(encode_value, *val), size, loc)
|
||||
ordering.mark_used(size, loc)
|
||||
@ -643,12 +664,19 @@ class BinaryEncoder:
|
||||
val = 1 if val else 0
|
||||
|
||||
# The size is built-in, emit it
|
||||
if size is None:
|
||||
raise Exception("Logic error, node size not set yet this is not an attribute!")
|
||||
|
||||
encode_value = f'>{enc}'
|
||||
self.__add_data(struct.pack(encode_value, val), size, loc)
|
||||
ordering.mark_used(size, loc)
|
||||
else:
|
||||
# Array value
|
||||
loc = ordering.get_next_int()
|
||||
if loc is None:
|
||||
raise BinaryEncodingException("Ran out of data when attempting allocate array location!")
|
||||
if size is None:
|
||||
raise Exception("Logic error, node size not set yet this is not an attribute!")
|
||||
|
||||
# The raw size in bytes
|
||||
elems = len(val)
|
||||
|
||||
@ -634,7 +634,7 @@ class Node:
|
||||
Returns:
|
||||
A string data type name. This string can be fed to typename_to_type to get the original type back.
|
||||
"""
|
||||
if self.__type is None:
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to fetch data type before setting type!')
|
||||
return self.__translated_type['name']
|
||||
|
||||
@ -647,7 +647,7 @@ class Node:
|
||||
Returns:
|
||||
An integer data length, or None if this node's element has variable length.
|
||||
"""
|
||||
if self.__type is None:
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to fetch data length before setting type!')
|
||||
if self.__translated_type['name'] in {'bin', 'str'}:
|
||||
return None
|
||||
@ -661,7 +661,7 @@ class Node:
|
||||
Returns:
|
||||
A character that can be passed to struct.pack or struct.unpack.
|
||||
"""
|
||||
if self.__type is None:
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to fetch data encoding before setting type!')
|
||||
return self.__translated_type['enc']
|
||||
|
||||
@ -781,6 +781,8 @@ class Node:
|
||||
Returns:
|
||||
True if this Node is a composite type, False otherwise.
|
||||
"""
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to fetch composite determination before setting type!')
|
||||
return self.__translated_type['composite']
|
||||
|
||||
def set_value(self, val: Any) -> None:
|
||||
@ -793,18 +795,22 @@ class Node:
|
||||
"""
|
||||
is_array = isinstance(val, (list, tuple))
|
||||
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to set value before setting type!')
|
||||
translated_type: Dict[str, Any] = self.__translated_type
|
||||
|
||||
# Handle composite types
|
||||
if self.__translated_type['composite']:
|
||||
if translated_type['composite']:
|
||||
if not is_array:
|
||||
raise NodeException('Input is not array, expected array')
|
||||
if len(val) != len(self.__translated_type['enc']):
|
||||
raise NodeException(f'Input array for {self.__translated_type["name"]} expected to be {len(self.__translated_type["enc"])} elements!')
|
||||
if len(val) != len(translated_type['enc']):
|
||||
raise NodeException(f'Input array for {translated_type["name"]} expected to be {len(translated_type["enc"])} elements!')
|
||||
is_array = False
|
||||
if is_array != self.__array:
|
||||
raise NodeException(f'Input {"is" if is_array else "is not"} array, expected {"array" if self.__array else "scalar"}')
|
||||
|
||||
def val_to_str(val: Any) -> Union[str, bytes]:
|
||||
if self.__translated_type['name'] == 'bool':
|
||||
if translated_type['name'] == 'bool':
|
||||
# Support user-built boolean types
|
||||
if val is True:
|
||||
return 'true'
|
||||
@ -813,9 +819,9 @@ class Node:
|
||||
|
||||
# Support construction from binary
|
||||
return 'true' if val != 0 else 'false'
|
||||
elif self.__translated_type['name'] == 'float':
|
||||
elif translated_type['name'] == 'float':
|
||||
return str(val)
|
||||
elif self.__translated_type['name'] == 'ip4':
|
||||
elif translated_type['name'] == 'ip4':
|
||||
try:
|
||||
# Support construction from binary
|
||||
ip = struct.unpack('BBBB', val)
|
||||
@ -827,13 +833,13 @@ class Node:
|
||||
return val
|
||||
|
||||
raise NodeException(f'Invalid value {val} for IP4 type')
|
||||
elif self.__translated_type['int']:
|
||||
elif translated_type['int']:
|
||||
return str(val)
|
||||
else:
|
||||
# This could return either a string or bytes.
|
||||
return val
|
||||
|
||||
if is_array or self.__translated_type['composite']:
|
||||
if is_array or translated_type['composite']:
|
||||
self.__value = [val_to_str(v) for v in val]
|
||||
else:
|
||||
self.__value = val_to_str(val)
|
||||
@ -846,23 +852,27 @@ class Node:
|
||||
Returns:
|
||||
A mixed value corresponding to this node's value. The returned value will be of the correct data type.
|
||||
"""
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to get value before setting type!')
|
||||
translated_type: Dict[str, Any] = self.__translated_type
|
||||
|
||||
def str_to_val(string: Union[str, bytes]) -> Any:
|
||||
if self.__translated_type['name'] == 'bool':
|
||||
if translated_type['name'] == 'bool':
|
||||
return True if string == 'true' else False
|
||||
elif self.__translated_type['name'] == 'float':
|
||||
elif translated_type['name'] == 'float':
|
||||
return float(string)
|
||||
elif self.__translated_type['name'] == 'ip4':
|
||||
elif translated_type['name'] == 'ip4':
|
||||
if not isinstance(string, str):
|
||||
raise Exception('Logic error, expected a string!')
|
||||
ip = [int(tup) for tup in string.split('.')]
|
||||
return struct.pack('BBBB', ip[0], ip[1], ip[2], ip[3])
|
||||
elif self.__translated_type['int']:
|
||||
elif translated_type['int']:
|
||||
return int(string)
|
||||
else:
|
||||
# At this point, we could be a string or bytes.
|
||||
return string
|
||||
|
||||
if self.__array or self.__translated_type['composite']:
|
||||
if self.__array or translated_type['composite']:
|
||||
return [str_to_val(v) for v in self.__value]
|
||||
else:
|
||||
return str_to_val(self.__value)
|
||||
@ -878,6 +888,10 @@ class Node:
|
||||
Returns:
|
||||
A string representing the XML-like data for this node and all children.
|
||||
"""
|
||||
if self.__translated_type is None:
|
||||
raise Exception('Logic error, tried to get XML representation before setting type!')
|
||||
translated_type: Dict[str, Any] = self.__translated_type
|
||||
|
||||
attrs_dict = copy.deepcopy(self.__attrs)
|
||||
order = sorted(attrs_dict.keys())
|
||||
if self.data_length != 0:
|
||||
@ -888,7 +902,7 @@ class Node:
|
||||
else:
|
||||
attrs_dict['__count'] = str(len(self.__value))
|
||||
order.insert(0, '__count')
|
||||
attrs_dict['__type'] = self.__translated_type['name']
|
||||
attrs_dict['__type'] = translated_type['name']
|
||||
order.insert(0, '__type')
|
||||
|
||||
def escape(val: Any, attr: bool=False) -> str:
|
||||
@ -912,14 +926,14 @@ class Node:
|
||||
attrs = ''
|
||||
|
||||
def get_val() -> str:
|
||||
if self.__array or self.__translated_type['composite']:
|
||||
if self.__array or translated_type['composite']:
|
||||
if self.__value is None:
|
||||
vals = ''
|
||||
else:
|
||||
vals = ' '.join([val for val in self.__value])
|
||||
elif self.__translated_type['name'] == 'str':
|
||||
elif translated_type['name'] == 'str':
|
||||
vals = escape(self.__value)
|
||||
elif self.__translated_type['name'] == 'bin':
|
||||
elif translated_type['name'] == 'bin':
|
||||
# Convert to a hex string
|
||||
def bin_to_hex(binary: int) -> str:
|
||||
val = hex(binary)[2:]
|
||||
|
||||
@ -66,7 +66,7 @@ class EAmuseProtocol:
|
||||
|
||||
return bytes(out)
|
||||
|
||||
def __decrypt(self, encryption_key: str, data: bytes) -> bytes:
|
||||
def __decrypt(self, encryption_key: Optional[str], data: bytes) -> bytes:
|
||||
"""
|
||||
Given data and an optional encryption key, decrypt the data.
|
||||
|
||||
@ -82,6 +82,7 @@ class EAmuseProtocol:
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
key: Optional[bytes] = None
|
||||
if encryption_key:
|
||||
# Key is concatenated with the shared secret above
|
||||
version, first, second = encryption_key.split('-')
|
||||
@ -91,8 +92,6 @@ class EAmuseProtocol:
|
||||
m = hashlib.md5()
|
||||
m.update(key)
|
||||
key = m.digest()
|
||||
else:
|
||||
key = None
|
||||
|
||||
if key:
|
||||
# This is an encrypted old-style packet
|
||||
@ -101,7 +100,7 @@ class EAmuseProtocol:
|
||||
# No encryption
|
||||
return data
|
||||
|
||||
def __encrypt(self, encryption_key: str, data: bytes) -> bytes:
|
||||
def __encrypt(self, encryption_key: Optional[str], data: bytes) -> bytes:
|
||||
"""
|
||||
Given data and an optional encryption key, encrypt the data.
|
||||
|
||||
@ -117,7 +116,7 @@ class EAmuseProtocol:
|
||||
# RC4 is symmetric
|
||||
return self.__decrypt(encryption_key, data)
|
||||
|
||||
def __decompress(self, compression: str, data: bytes) -> bytes:
|
||||
def __decompress(self, compression: Optional[str], data: bytes) -> bytes:
|
||||
"""
|
||||
Given data and an optional compression scheme, decompress the data.
|
||||
|
||||
@ -133,21 +132,17 @@ class EAmuseProtocol:
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
if compression:
|
||||
if compression is None or compression == 'none':
|
||||
# This isn't compressed
|
||||
return data
|
||||
elif compression == 'lz77':
|
||||
# This is a compressed new-style packet
|
||||
lz = Lz77()
|
||||
return lz.decompress(data)
|
||||
else:
|
||||
raise EAmuseException(f'Unknown compression {compression}')
|
||||
if compression is None or compression == 'none':
|
||||
# This isn't compressed
|
||||
return data
|
||||
elif compression == 'lz77':
|
||||
# This is a compressed new-style packet
|
||||
lz = Lz77()
|
||||
return lz.decompress(data)
|
||||
else:
|
||||
raise EAmuseException(f'Unknown compression {compression}')
|
||||
|
||||
# No compression
|
||||
return data
|
||||
|
||||
def __compress(self, compression: str, data: bytes) -> bytes:
|
||||
def __compress(self, compression: Optional[str], data: bytes) -> bytes:
|
||||
"""
|
||||
Given data and an optional compression scheme, compress the data.
|
||||
|
||||
@ -163,19 +158,15 @@ class EAmuseProtocol:
|
||||
if data is None:
|
||||
return None
|
||||
|
||||
if compression:
|
||||
if compression is None or compression == 'none':
|
||||
# This isn't compressed
|
||||
return data
|
||||
elif compression == 'lz77':
|
||||
# This is a compressed new-style packet
|
||||
lz = Lz77()
|
||||
return lz.compress(data)
|
||||
else:
|
||||
raise EAmuseException(f'Unknown compression {compression}')
|
||||
|
||||
# No compression
|
||||
return data
|
||||
if compression is None or compression == 'none':
|
||||
# This isn't compressed
|
||||
return data
|
||||
elif compression == 'lz77':
|
||||
# This is a compressed new-style packet
|
||||
lz = Lz77()
|
||||
return lz.compress(data)
|
||||
else:
|
||||
raise EAmuseException(f'Unknown compression {compression}')
|
||||
|
||||
def __decode(self, data: bytes) -> Node:
|
||||
"""
|
||||
|
||||
@ -48,7 +48,7 @@ class InputStream:
|
||||
return bytedata
|
||||
return None
|
||||
|
||||
def read_byte(self) -> bytes:
|
||||
def read_byte(self) -> Optional[bytes]:
|
||||
"""
|
||||
Grab the next byte at the current position. If no byte is available,
|
||||
return None.
|
||||
@ -58,7 +58,7 @@ class InputStream:
|
||||
"""
|
||||
return self.read_blob(1)
|
||||
|
||||
def read_int(self, size: int=1, is_unsigned: bool=True) -> int:
|
||||
def read_int(self, size: int=1, is_unsigned: bool=True) -> Optional[int]:
|
||||
"""
|
||||
Grab the next integer of size 'size' at the current position. If not enough
|
||||
bytes are available to decode this integer, return None.
|
||||
|
||||
@ -69,7 +69,7 @@ class XmlDecoder:
|
||||
# Skip these, handled
|
||||
continue
|
||||
else:
|
||||
node.set_attribute(attr, attributes.get(attr))
|
||||
node.set_attribute(attr, attributes[attr])
|
||||
|
||||
self.current.append(node)
|
||||
|
||||
@ -298,7 +298,7 @@ class XmlDecoder:
|
||||
if empty:
|
||||
self.__end_element(tag)
|
||||
|
||||
def get_tree(self) -> Node:
|
||||
def get_tree(self) -> Optional[Node]:
|
||||
"""
|
||||
Walk the XML document and parse into nodes.
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user