Make use of yield in __iter__

This commit is contained in:
Emily 2023-03-15 03:33:24 +01:00
parent 73eaa0e89b
commit 404f71dddd
2 changed files with 3 additions and 10 deletions

View File

@ -167,15 +167,8 @@ class Bitstream:
return self return self
def __iter__(self) -> Bitstream: def __iter__(self) -> Bitstream:
self._iter_index = 0 for i in range(len(self)):
return self yield self[i]
def __next__(self) -> bool:
if self._iter_index >= self._bit_length:
raise StopIteration
bit = self[self._iter_index]
self._iter_index += 1
return bit
def __repr__(self) -> str: def __repr__(self) -> str:
return "<Bitstream addr:0x%012x offset:%d len:%d data:'%s'>" % ( return "<Bitstream addr:0x%012x offset:%d len:%d data:'%s'>" % (

View File

@ -181,7 +181,7 @@ def test_add(
]) ])
def test_iter(init_values: List[Union[bool, List[bool], bytes, str, Bitstream]]): def test_iter(init_values: List[Union[bool, List[bool], bytes, str, Bitstream]]):
bitstream = Bitstream(*init_values) bitstream = Bitstream(*init_values)
assert len(iter(bitstream)) == len(bitstream) assert len(list(bitstream)) == len(bitstream)
for index, bit in enumerate(bitstream): for index, bit in enumerate(bitstream):
assert bit == bitstream[index] assert bit == bitstream[index]