dvc/dvc/utils/stream.py

import io


class IterStream(io.RawIOBase):
    """Wraps an iterator yielding bytes as a file object"""

    def __init__(self, iterator):  # pylint: disable=super-init-not-called
        self.iterator = iterator
        self.leftover = None

    def readable(self):
        return True

    # Python 3 requires only .readinto() method, it still uses other ones
    # under some circumstances and falls back if those are absent. Since
    # iterator already constructs byte strings for us, .readinto() is not the
    # most optimal, so we provide .read1() too.

    def readinto(self, b):
        try:
            n = len(b)  # We're supposed to return at most this much
            chunk = self.leftover or next(self.iterator)
            output, self.leftover = chunk[:n], chunk[n:]

            n_out = len(output)
            b[:n_out] = output
            return n_out
        except StopIteration:
            return 0  # indicate EOF

    readinto1 = readinto

    def read1(self, n=-1):
        try:
            chunk = self.leftover or next(self.iterator)
        except StopIteration:
            return b""

        # Return an arbitrary number or bytes
        if n <= 0:
            self.leftover = None
            return chunk

        output, self.leftover = chunk[:n], chunk[n:]
        return output