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