add type hints to common.py
[~helmut/crossqa.git] / common.py
1 import typing
2
3 ByteIterable = typing.Iterable[bytes]
4
5 def yield_lines(iterable: ByteIterable) -> ByteIterable:
6     """Converts an arbitrary bytes iterable into an iterable that yields whole
7     lines. The final byte of each returned value (except possibly the last one)
8     is a newline or carriage return character. The concatenation of the input
9     iterable equals the concatenation of the output iterable."""
10     buff = b""
11     for data in iterable:
12         buff += data
13         parts = buff.splitlines(True)
14         buff = parts.pop()
15         yield from parts
16     if buff:
17         yield buff
18
19 def decompress_stream(iterable: ByteIterable, decompressor) -> ByteIterable:
20     """Decompress an iterable of bytes using the given decompressor into
21     another (decompressed) iterable of bytes. The decompressor can be a
22     bz2.BZ2Decompressor or lzma.LZMADecompressor instance."""
23     for data in iterable:
24         data = decompressor.decompress(data)
25         yield data
26     if hasattr(decompressor, "flush"):
27         yield decompressor.flush()
28
29 def yield_chunks(filelike, chunksize=65536) -> ByteIterable:
30     """Read the given file in chunks of the given size. Returns an itrable
31     of contents. If the file is binary, it yields bytes, otherwise str."""
32     while True:
33         data = filelike.read(chunksize)
34         if not data:
35             break
36         yield data