U
    Mh                     @   sR   d dl mZ d dlmZ d dlmZ dgZedG dd deeeef  Z	dS )    )Tuple)functional_datapipe)IterDataPipeStreamReaderIterDataPipeZread_from_streamc                   @   s"   e Zd ZdZdddZdd ZdS )r   a  
    Given IO streams and their label names, yield bytes with label name as 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 until 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')]
    Nc                 C   s   || _ || _d S N)datapipechunk)selfr   r    r
   ^/var/www/html/venv/lib/python3.8/site-packages/torch/utils/data/datapipes/iter/streamreader.py__init__   s    z!StreamReaderIterDataPipe.__init__c                 c   s:   | j D ].\}}|| j}|s(|  q||fV  qqd S r   )r   readr   close)r	   Zfurlstreamdr
   r
   r   __iter__"   s    z!StreamReaderIterDataPipe.__iter__)N)__name__
__module____qualname____doc__r   r   r
   r
   r
   r   r   	   s   
N)
typingr   Z%torch.utils.data.datapipes._decoratorr   Z#torch.utils.data.datapipes.datapiper   __all__strbytesr   r
   r
   r
   r   <module>   s
   