bazarr/libs/requests_toolbelt/downloadutils/tee.py

"""Tee function implementations."""
import io

_DEFAULT_CHUNKSIZE = 65536

__all__ = ['tee', 'tee_to_file', 'tee_to_bytearray']


def _tee(response, callback, chunksize, decode_content):
    for chunk in response.raw.stream(amt=chunksize,
                                     decode_content=decode_content):
        callback(chunk)
        yield chunk


def tee(response, fileobject, chunksize=_DEFAULT_CHUNKSIZE,
        decode_content=None):
    """Stream the response both to the generator and a file.

    This will stream the response body while writing the bytes to
    ``fileobject``.

    Example usage:

    .. code-block:: python

        resp = requests.get(url, stream=True)
        with open('save_file', 'wb') as save_file:
            for chunk in tee(resp, save_file):
                # do stuff with chunk

    .. code-block:: python

        import io

        resp = requests.get(url, stream=True)
        fileobject = io.BytesIO()

        for chunk in tee(resp, fileobject):
            # do stuff with chunk

    :param response: Response from requests.
    :type response: requests.Response
    :param fileobject: Writable file-like object.
    :type fileobject: file, io.BytesIO
    :param int chunksize: (optional), Size of chunk to attempt to stream.
    :param bool decode_content: (optional), If True, this will decode the
        compressed content of the response.
    :raises: TypeError if the fileobject wasn't opened with the right mode
        or isn't a BytesIO object.
    """
    # We will be streaming the raw bytes from over the wire, so we need to
    # ensure that writing to the fileobject will preserve those bytes. On
    # Python3, if the user passes an io.StringIO, this will fail, so we need
    # to check for BytesIO instead.
    if not ('b' in getattr(fileobject, 'mode', '') or
            isinstance(fileobject, io.BytesIO)):
        raise TypeError('tee() will write bytes directly to this fileobject'
                        ', it must be opened with the "b" flag if it is a file'
                        ' or inherit from io.BytesIO.')

    return _tee(response, fileobject.write, chunksize, decode_content)


def tee_to_file(response, filename, chunksize=_DEFAULT_CHUNKSIZE,
                decode_content=None):
    """Stream the response both to the generator and a file.

    This will open a file named ``filename`` and stream the response body
    while writing the bytes to the opened file object.

    Example usage:

    .. code-block:: python

        resp = requests.get(url, stream=True)
        for chunk in tee_to_file(resp, 'save_file'):
            # do stuff with chunk

    :param response: Response from requests.
    :type response: requests.Response
    :param str filename: Name of file in which we write the response content.
    :param int chunksize: (optional), Size of chunk to attempt to stream.
    :param bool decode_content: (optional), If True, this will decode the
        compressed content of the response.
    """
    with open(filename, 'wb') as fd:
        for chunk in tee(response, fd, chunksize, decode_content):
            yield chunk


def tee_to_bytearray(response, bytearr, chunksize=_DEFAULT_CHUNKSIZE,
                     decode_content=None):
    """Stream the response both to the generator and a bytearray.

    This will stream the response provided to the function, add them to the
    provided :class:`bytearray` and yield them to the user.

    .. note::

        This uses the :meth:`bytearray.extend` by default instead of passing
        the bytearray into the ``readinto`` method.

    Example usage:

    .. code-block:: python

        b = bytearray()
        resp = requests.get(url, stream=True)
        for chunk in tee_to_bytearray(resp, b):
            # do stuff with chunk

    :param response: Response from requests.
    :type response: requests.Response
    :param bytearray bytearr: Array to add the streamed bytes to.
    :param int chunksize: (optional), Size of chunk to attempt to stream.
    :param bool decode_content: (optional), If True, this will decode the
        compressed content of the response.
    """
    if not isinstance(bytearr, bytearray):
        raise TypeError('tee_to_bytearray() expects bytearr to be a '
                        'bytearray')
    return _tee(response, bytearr.extend, chunksize, decode_content)