1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
|
import io
from typing import IO, AnyStr
def read_until_size_or_end(stream: IO[AnyStr], size: int) -> io.BytesIO:
pos = 0
result = io.BytesIO()
while True:
try:
returned = stream.read(size - pos)
pos += len(returned)
result.write(returned)
if len(returned) == 0 or pos == size:
result.seek(0, io.SEEK_SET)
return result
except BlockingIOError:
continue
class ChainStream(io.RawIOBase):
"""
https://stackoverflow.com/questions/24528278/stream-multiple-files-into-a-readable-object-in-python
"""
def __init__(self, streams):
self.leftover = b""
self.stream_iter = iter(streams)
try:
self.stream = next(self.stream_iter)
except StopIteration:
self.stream = None
def readable(self):
return True
def _read_next_chunk(self, max_length):
# Return 0 or more bytes from the current stream, first returning all
# leftover bytes. If the stream is closed returns b''
if self.leftover:
return self.leftover
elif self.stream is not None:
return self.stream.read(max_length)
else:
return b""
def readinto(self, b):
buffer_length = len(b)
chunk = self._read_next_chunk(buffer_length)
while len(chunk) == 0:
# move to next stream
if self.stream is not None:
self.stream.close()
try:
self.stream = next(self.stream_iter)
chunk = self._read_next_chunk(buffer_length)
except StopIteration:
# No more streams to chain together
self.stream = None
return 0 # indicate EOF
output, self.leftover = chunk[:buffer_length], chunk[buffer_length:]
b[: len(output)] = output
return len(output)
def chain_streams(streams, buffer_size=io.DEFAULT_BUFFER_SIZE):
"""
Chain an iterable of streams together into a single buffered stream.
Usage:
def generate_open_file_streams():
for file in filenames:
yield open(file, 'rb')
f = chain_streams(generate_open_file_streams())
f.read()
"""
return io.BufferedReader(ChainStream(streams), buffer_size=buffer_size)
|
