diff options
| -rw-r--r-- | discord/oggparse.py | 33 |
1 files changed, 22 insertions, 11 deletions
diff --git a/discord/oggparse.py b/discord/oggparse.py index f2a14927..e0347d2c 100644 --- a/discord/oggparse.py +++ b/discord/oggparse.py @@ -22,8 +22,12 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from __future__ import annotations + import struct +from typing import TYPE_CHECKING, ClassVar, IO, Generator, Tuple, Optional + from .errors import DiscordException __all__ = ( @@ -40,22 +44,29 @@ class OggError(DiscordException): # https://tools.ietf.org/html/rfc7845 class OggPage: - _header = struct.Struct('<xBQIIIB') - - def __init__(self, stream): + _header: ClassVar[struct.Struct] = struct.Struct('<xBQIIIB') + if TYPE_CHECKING: + flag: int + gran_pos: int + serial: int + pagenum: int + crc: int + segnum: int + + def __init__(self, stream: IO[bytes]) -> None: try: header = stream.read(struct.calcsize(self._header.format)) self.flag, self.gran_pos, self.serial, \ self.pagenum, self.crc, self.segnum = self._header.unpack(header) - self.segtable = stream.read(self.segnum) + self.segtable: bytes = stream.read(self.segnum) bodylen = sum(struct.unpack('B'*self.segnum, self.segtable)) - self.data = stream.read(bodylen) + self.data: bytes = stream.read(bodylen) except Exception: raise OggError('bad data stream') from None - def iter_packets(self): + def iter_packets(self) -> Generator[Tuple[bytes, bool], None, None]: packetlen = offset = 0 partial = True @@ -74,10 +85,10 @@ class OggPage: yield self.data[offset:], False class OggStream: - def __init__(self, stream): - self.stream = stream + def __init__(self, stream: IO[bytes]) -> None: + self.stream: IO[bytes] = stream - def _next_page(self): + def _next_page(self) -> Optional[OggPage]: head = self.stream.read(4) if head == b'OggS': return OggPage(self.stream) @@ -86,13 +97,13 @@ class OggStream: else: raise OggError('invalid header magic') - def _iter_pages(self): + def _iter_pages(self) -> Generator[OggPage, None, None]: page = self._next_page() while page: yield page page = self._next_page() - def iter_packets(self): + def iter_packets(self) -> Generator[bytes, None, None]: partial = b'' for page in self._iter_pages(): for data, complete in page.iter_packets(): |