|
|
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) |
|
|
|
