Sunpy/sampy3
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Removed str type from bitstream (str is ambiguous)

Browse Source
This commit is contained in:
Emily 2023-03-15 03:43:18 +01:00
parent 404f71dddd
commit 446ab11a61
2 changed files with 36 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
sampy/raknet/bitstream.py
View File

@ -8,12 +8,12 @@ SEEK_CUR = 1
SEEK_END = 2 SEEK_END = 2
def get_bit_length(value: Union[bool, List[bool], bytes, str, Bitstream]) -> int: def get_bit_length(value: Union[bool, List[bool], bytes, Bitstream]) -> int:
if type(value) is bool: if type(value) is bool:
return 1 return 1
elif type(value) is list and all(type(v) is bool for v in value): elif type(value) is list and all(type(v) is bool for v in value):
return len(value) return len(value)
elif type(value) in (bytes, str): elif type(value) is bytes:
return len(value) << 3 return len(value) << 3
elif type(value) is Bitstream: elif type(value) is Bitstream:
return len(value) return len(value)
@ -36,7 +36,7 @@ def int_to_bits(number: int, bit_length: int, little_endian: bool = False) -> Li
class Bitstream: class Bitstream:
def __init__(self, *values: Union[bool, List[bool], bytes, str, Bitstream]): def __init__(self, *values: Union[bool, List[bool], bytes, Bitstream]):
self._bytearray = bytearray() self._bytearray = bytearray()
self._bit_length = 0 self._bit_length = 0
self._offset = 0 self._offset = 0
@ -66,9 +66,6 @@ class Bitstream:
self._bytearray.pop() self._bytearray.pop()
return bytes(self._bytearray) return bytes(self._bytearray)
def __str__(self) -> str:
return bytes(self).decode()
def __int__(self) -> int: def __int__(self) -> int:
num = 0 num = 0
for bit in self: for bit in self:
@ -85,13 +82,10 @@ class Bitstream:
mask = 1 << (7 - (bit_index % 8)) mask = 1 << (7 - (bit_index % 8))
return bool(self._bytearray[bit_index >> 3] & mask) return bool(self._bytearray[bit_index >> 3] & mask)
def __setitem__(self, bit_index: int, value: Union[bool, List[bool], bytes, str, Bitstream]): def __setitem__(self, bit_index: int, value: Union[bool, List[bool], bytes, Bitstream]):
if bit_index >= self._bit_length: if bit_index >= self._bit_length:
raise IndexError("bit index out of range") raise IndexError("bit index out of range")
if type(value) is str:
value = value.encode()
if type(value) is bytes: if type(value) is bytes:
self[bit_index] = Bitstream.from_bytes(value) self[bit_index] = Bitstream.from_bytes(value)
return return
@ -111,13 +105,11 @@ class Bitstream:
for i in range(value_bit_length): for i in range(value_bit_length):
self[bit_index + i] = value[i] self[bit_index + i] = value[i]
else: else:
raise TypeError("Expected bool, list[bool], bytes, str or Bitstream") raise TypeError("Expected bool, list[bool], bytes or Bitstream")
def __eq__(self, other: Any) -> bool: def __eq__(self, other: Any) -> bool:
if type(other) is bool: if type(other) is bool:
other = [other] other = [other]
if type(other) is str:
other = other.encode()
if isinstance(other, Bitstream): if isinstance(other, Bitstream):
pass pass
@ -126,7 +118,7 @@ class Bitstream:
elif type(other) is list and all(type(v) is bool for v in other): elif type(other) is list and all(type(v) is bool for v in other):
pass pass
else: else:
raise TypeError("Expected Bitstream, bytes, str or list[bool]") raise TypeError("Expected Bitstream, bytes or list[bool]")
if isinstance(other, Bitstream) or type(other) is list: if isinstance(other, Bitstream) or type(other) is list:
if len(self) != len(other): if len(self) != len(other):
@ -151,7 +143,7 @@ class Bitstream:
self.extend(copy) self.extend(copy)
return self return self
def __add__(self, value: Union[bool, List[bool], bytes, str, Bitstream]) -> Bitstream: def __add__(self, value: Union[bool, List[bool], bytes, Bitstream]) -> Bitstream:
copy = self.copy() copy = self.copy()
if type(value) is bool: if type(value) is bool:
copy.append(value) copy.append(value)
@ -159,7 +151,7 @@ class Bitstream:
copy.extend(value) copy.extend(value)
return copy return copy
def __iadd__(self, value: Union[bool, List[bool], bytes, str, Bitstream]) -> Bitstream: def __iadd__(self, value: Union[bool, List[bool], bytes, Bitstream]) -> Bitstream:
if type(value) is bool: if type(value) is bool:
self.append(value) self.append(value)
else: else:
@ -196,15 +188,15 @@ class Bitstream:
self._bit_length += 1 self._bit_length += 1
self[prev_bit_length] = value self[prev_bit_length] = value
def extend(self, value: Union[List[bool], bytes, str, Bitstream]): def extend(self, value: Union[List[bool], bytes, Bitstream]):
if type(value) in (bytes, str): if type(value) is bytes:
pass pass
elif isinstance(value, Bitstream): elif isinstance(value, Bitstream):
pass pass
elif type(value) is list and all(type(v) is bool for v in value): elif type(value) is list and all(type(v) is bool for v in value):
pass pass
else: else:
raise TypeError("Expected list[bool], bytes, str or Bitstream") raise TypeError("Expected list[bool], bytes or Bitstream")
value_bit_length = get_bit_length(value) value_bit_length = get_bit_length(value)
if value_bit_length == 0: if value_bit_length == 0:
@ -229,10 +221,10 @@ class Bitstream:
def read( def read(
self, self,
type: Type[Union[bool, bytes, str, Bitstream, int]], type: Type[Union[bool, bytes, Bitstream, int]],
bit_length: int, bit_length: int,
bit_index: Optional[int] = None bit_index: Optional[int] = None
) -> Union[List[bool], bytes, str, Bitstream, int]: ) -> Union[List[bool], bytes, Bitstream, int]:
start = self._offset if bit_index is None else bit_index start = self._offset if bit_index is None else bit_index
if (start + bit_length) > self._bit_length or (start < 0): if (start + bit_length) > self._bit_length or (start < 0):
@ -251,14 +243,12 @@ class Bitstream:
return bitstream return bitstream
elif type is bytes: elif type is bytes:
return bytes(bitstream) return bytes(bitstream)
elif type is str:
return str(bitstream)
elif type is int: elif type is int:
return int(bitstream) return int(bitstream)
raise TypeError("Invalid data type") raise TypeError("Invalid data type")
def write(self, value: Union[bool, List[bool], bytes, str, Bitstream], bit_index: Optional[int] = None): def write(self, value: Union[bool, List[bool], bytes, Bitstream], bit_index: Optional[int] = None):
start = self._offset if bit_index is None else bit_index start = self._offset if bit_index is None else bit_index
if start < 0: if start < 0:

55
tests/test_bitstream.py
View File

@ -38,11 +38,11 @@ def test_from_int(value: int, length: int, expected: int):
[True, False, False, True, True], [True, False, False, True, True],
[[True], [False, True], [False, False]], [[True], [False, True], [False, False]],
[b"a", b"bc"], [b"a", b"bc"],
["aa", "b"], [b"aa", b"b"],
[Bitstream.from_bytes(b"A"), Bitstream.from_int(0b1011, 4), Bitstream.from_bytes(b"B")], [Bitstream.from_bytes(b"A"), Bitstream.from_int(0b1011, 4), Bitstream.from_bytes(b"B")],
[True, b"A", "B", [False, True], Bitstream.from_int(0b10110, 5), False], [True, b"A", b"B", [False, True], Bitstream.from_int(0b10110, 5), False],
]) ])
def test_init(values: List[Union[bool, List[bool], bytes, str, Bitstream]]): def test_init(values: List[Union[bool, List[bool], bytes, Bitstream]]):
bitstream = Bitstream(*values) bitstream = Bitstream(*values)
bitstream.seek(0) bitstream.seek(0)
@ -79,10 +79,9 @@ def test_getitem_index_error():
True, True,
[True, True], [True, True],
b"A", b"A",
"B",
Bitstream(True, b"C"), Bitstream(True, b"C"),
]) ])
def test_setitem(value: Union[bool, List[bool], bytes, str, Bitstream]): def test_setitem(value: Union[bool, List[bool], bytes, Bitstream]):
bit_length = get_bit_length(value) bit_length = get_bit_length(value)
read_type = bool if type(value) is list else type(value) read_type = bool if type(value) is list else type(value)
for length in range(10): for length in range(10):
@ -111,13 +110,8 @@ def test_setitem(value: Union[bool, List[bool], bytes, str, Bitstream]):
), ),
( (
b"A", b"A",
[b"A", "A", Bitstream.from_bytes(b"A"), Bitstream.from_bytes(b"A").bits], [b"A", b"A", Bitstream.from_bytes(b"A"), Bitstream.from_bytes(b"A").bits],
[[], [False], b"B", "B", Bitstream.from_bytes(b"B"), Bitstream.from_bytes(b"B").bits] [[], [False], b"B", b"B", Bitstream.from_bytes(b"B"), Bitstream.from_bytes(b"B").bits]
),
(
"B",
[b"B", "B", Bitstream.from_bytes(b"B"), Bitstream.from_bytes(b"B").bits],
[[], [False], b"A", "A", Bitstream.from_bytes(b"A"), Bitstream.from_bytes(b"A").bits]
), ),
( (
b"ABC", b"ABC",
@ -126,9 +120,9 @@ def test_setitem(value: Union[bool, List[bool], bytes, str, Bitstream]):
), ),
]) ])
def test_eq( def test_eq(
value: Union[bool, List[bool], bytes, str, Bitstream], value: Union[bool, List[bool], bytes, Bitstream],
equals: List[Union[bool, List[bool], bytes, str, Bitstream]], equals: List[Union[bool, List[bool], bytes, Bitstream]],
not_equals: List[Union[bool, List[bool], bytes, str, Bitstream]]): not_equals: List[Union[bool, List[bool], bytes, Bitstream]]):
bitstream = Bitstream(value) bitstream = Bitstream(value)
for equal in equals: for equal in equals:
assert bitstream == equal assert bitstream == equal
@ -143,9 +137,9 @@ def test_eq(
([False], 3, [False, False, False]), ([False], 3, [False, False, False]),
]) ])
def test_mul( def test_mul(
values: Union[bool, List[bool], bytes, str, Bitstream], values: Union[bool, List[bool], bytes, Bitstream],
number: int, number: int,
expected: Union[List[bool], bytes, str, Bitstream]): expected: Union[List[bool], bytes, Bitstream]):
bitstream = Bitstream(values) bitstream = Bitstream(values)
assert (bitstream * number) == expected assert (bitstream * number) == expected
@ -161,9 +155,9 @@ def test_mul(
(b"A", b"B", b"AB"), (b"A", b"B", b"AB"),
]) ])
def test_add( def test_add(
init_values: Union[bool, List[bool], bytes, str, Bitstream], init_values: Union[bool, List[bool], bytes, Bitstream],
add_value: Union[bool, List[bool], bytes, str, Bitstream], add_value: Union[bool, List[bool], bytes, Bitstream],
expected: Union[List[bool], bytes, str, Bitstream]): expected: Union[List[bool], bytes, Bitstream]):
bitstream = Bitstream(init_values) bitstream = Bitstream(init_values)
assert (bitstream + add_value) == expected assert (bitstream + add_value) == expected
@ -179,7 +173,7 @@ def test_add(
[b"A"], [b"A"],
[True, b"B"], [True, b"B"],
]) ])
def test_iter(init_values: List[Union[bool, List[bool], bytes, str, Bitstream]]): def test_iter(init_values: List[Union[bool, List[bool], bytes, Bitstream]]):
bitstream = Bitstream(*init_values) bitstream = Bitstream(*init_values)
assert len(list(bitstream)) == len(bitstream) assert len(list(bitstream)) == len(bitstream)
for index, bit in enumerate(bitstream): for index, bit in enumerate(bitstream):
@ -194,7 +188,7 @@ def test_iter(init_values: List[Union[bool, List[bool], bytes, str, Bitstream]])
([[False] * 7, True], b"\x01"), ([[False] * 7, True], b"\x01"),
]) ])
def test_bytes( def test_bytes(
init_values: List[Union[bool, List[bool], bytes, str, Bitstream]], init_values: List[Union[bool, List[bool], bytes, Bitstream]],
expected: bytes): expected: bytes):
assert bytes(Bitstream(*init_values)) == expected assert bytes(Bitstream(*init_values)) == expected
@ -207,7 +201,7 @@ def test_bytes(
([False] * 7 + [True], [False] * 7 + [True]), ([False] * 7 + [True], [False] * 7 + [True]),
]) ])
def test_bits( def test_bits(
init_values: List[Union[bool, List[bool], bytes, str, Bitstream]], init_values: List[Union[bool, List[bool], bytes, Bitstream]],
expected: List[bool]): expected: List[bool]):
assert Bitstream(*init_values).bits == expected assert Bitstream(*init_values).bits == expected
@ -236,10 +230,10 @@ def test_append(lst: List[bool]):
@pytest.mark.parametrize("extend_value", [ @pytest.mark.parametrize("extend_value", [
[True, False, True], [True, False, True],
b"A", b"A",
"B", b"BB",
Bitstream(True, b"C"), Bitstream(True, b"C"),
]) ])
def test_extend(extend_value: Union[List[bool], bytes, str, Bitstream]): def test_extend(extend_value: Union[List[bool], bytes, Bitstream]):
for i in range(10): for i in range(10):
bitstream = Bitstream([False] * i) bitstream = Bitstream([False] * i)
bitstream.extend(extend_value) bitstream.extend(extend_value)
@ -304,11 +298,6 @@ def test_read():
bytes, 8, 1, bytes, 8, 1,
b"A" b"A"
), ),
(
Bitstream(True, b"A", False, True, False),
str, 8, 1,
"A"
),
( (
Bitstream(True, b"A", False, True, False), Bitstream(True, b"A", False, True, False),
Bitstream, 1 + 8 + 3, 0, Bitstream, 1 + 8 + 3, 0,
@ -322,10 +311,10 @@ def test_read():
]) ])
def test_read_index( def test_read_index(
bitstream: Bitstream, bitstream: Bitstream,
read_type: Type[Union[bool, bytes, str, Bitstream, int]], read_type: Type[Union[bool, bytes, Bitstream, int]],
bit_length: int, bit_length: int,
bit_index: Optional[int], bit_index: Optional[int],
expected: Union[List[bool], bytes, str, Bitstream, int]): expected: Union[List[bool], bytes, Bitstream, int]):
value = bitstream.read(read_type, bit_length, bit_index) value = bitstream.read(read_type, bit_length, bit_index)
if read_type is bool: if read_type is bool:
@ -360,7 +349,7 @@ def test_read_error():
(Bitstream(False), True, 0, 0), (Bitstream(False), True, 0, 0),
(Bitstream(False), True, None, 0), (Bitstream(False), True, None, 0),
]) ])
def test_write(bitstream: Bitstream, value: Union[bool, List[bool], bytes, str, Bitstream], bit_index: Optional[int], current_offset: Optional[int]): def test_write(bitstream: Bitstream, value: Union[bool, List[bool], bytes, Bitstream], bit_index: Optional[int], current_offset: Optional[int]):
if current_offset is not None: if current_offset is not None:
bitstream.seek(current_offset) bitstream.seek(current_offset)