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
|
from typing import Tuple
from torch.utils.data.datapipes._decorator import functional_datapipe
from torch.utils.data.datapipes.datapipe import IterDataPipe
__all__ = ["StreamReaderIterDataPipe", ]
@functional_datapipe('read_from_stream')
class StreamReaderIterDataPipe(IterDataPipe[Tuple[str, bytes]]):
r"""
Given IO streams and their label names, yields bytes with label
name in a tuple (functional name: ``read_from_stream``).
Args:
datapipe: Iterable DataPipe provides label/URL and byte stream
chunk: Number of bytes to be read from stream per iteration.
If ``None``, all bytes will be read util the EOF.
Example:
>>> # xdoctest: +SKIP
>>> from torchdata.datapipes.iter import IterableWrapper, StreamReader
>>> from io import StringIO
>>> dp = IterableWrapper([("alphabet", StringIO("abcde"))])
>>> list(StreamReader(dp, chunk=1))
[('alphabet', 'a'), ('alphabet', 'b'), ('alphabet', 'c'), ('alphabet', 'd'), ('alphabet', 'e')]
"""
def __init__(self, datapipe, chunk=None):
self.datapipe = datapipe
self.chunk = chunk
def __iter__(self):
for furl, stream in self.datapipe:
while True:
d = stream.read(self.chunk)
if not d:
stream.close()
break
yield (furl, d)
|
