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