codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2"""Utilities for dealing with streamed requests."""3import os.path4import re5 6from .. import exceptions as exc7 8# Regular expressions stolen from werkzeug/http.py9# cd2c97bb0a076da2322f11adce0b2731f9193396 L62-L6410_QUOTED_STRING_RE = r'"[^"\\]*(?:\\.[^"\\]*)*"'11_OPTION_HEADER_PIECE_RE = re.compile(12    r';\s*(%s|[^\s;=]+)\s*(?:=\s*(%s|[^;]+))?\s*' % (_QUOTED_STRING_RE,13                                                     _QUOTED_STRING_RE)14)15_DEFAULT_CHUNKSIZE = 51216 17 18def _get_filename(content_disposition):19    for match in _OPTION_HEADER_PIECE_RE.finditer(content_disposition):20        k, v = match.groups()21        if k == 'filename':22            # ignore any directory paths in the filename23            return os.path.split(v)[1]24    return None25 26 27def get_download_file_path(response, path):28    """29    Given a response and a path, return a file path for a download.30 31    If a ``path`` parameter is a directory, this function will parse the32    ``Content-Disposition`` header on the response to determine the name of the33    file as reported by the server, and return a file path in the specified34    directory.35 36    If ``path`` is empty or None, this function will return a path relative37    to the process' current working directory.38 39    If path is a full file path, return it.40 41    :param response: A Response object from requests42    :type response: requests.models.Response43    :param str path: Directory or file path.44    :returns: full file path to download as45    :rtype: str46    :raises: :class:`requests_toolbelt.exceptions.StreamingError`47    """48    path_is_dir = path and os.path.isdir(path)49 50    if path and not path_is_dir:51        # fully qualified file path52        filepath = path53    else:54        response_filename = _get_filename(55            response.headers.get('content-disposition', '')56        )57        if not response_filename:58            raise exc.StreamingError('No filename given to stream response to')59 60        if path_is_dir:61            # directory to download to62            filepath = os.path.join(path, response_filename)63        else:64            # fallback to downloading to current working directory65            filepath = response_filename66 67    return filepath68 69 70def stream_response_to_file(response, path=None, chunksize=_DEFAULT_CHUNKSIZE):71    """Stream a response body to the specified file.72 73    Either use the ``path`` provided or use the name provided in the74    ``Content-Disposition`` header.75 76    .. warning::77 78        If you pass this function an open file-like object as the ``path``79        parameter, the function will not close that file for you.80 81    .. warning::82 83        This function will not automatically close the response object84        passed in as the ``response`` parameter.85 86    If a ``path`` parameter is a directory, this function will parse the87    ``Content-Disposition`` header on the response to determine the name of the88    file as reported by the server, and return a file path in the specified89    directory. If no ``path`` parameter is supplied, this function will default90    to the process' current working directory.91 92    .. code-block:: python93 94        import requests95        from requests_toolbelt import exceptions96        from requests_toolbelt.downloadutils import stream97 98        r = requests.get(url, stream=True)99        try:100            filename = stream.stream_response_to_file(r)101        except exceptions.StreamingError as e:102            # The toolbelt could not find the filename in the103            # Content-Disposition104            print(e.message)105 106    You can also specify the filename as a string. This will be passed to107    the built-in :func:`open` and we will read the content into the file.108 109    .. code-block:: python110 111        import requests112        from requests_toolbelt.downloadutils import stream113 114        r = requests.get(url, stream=True)115        filename = stream.stream_response_to_file(r, path='myfile')116 117    If the calculated download file path already exists, this function will118    raise a StreamingError.119 120    Instead, if you want to manage the file object yourself, you need to121    provide either a :class:`io.BytesIO` object or a file opened with the122    `'b'` flag. See the two examples below for more details.123 124    .. code-block:: python125 126        import requests127        from requests_toolbelt.downloadutils import stream128 129        with open('myfile', 'wb') as fd:130            r = requests.get(url, stream=True)131            filename = stream.stream_response_to_file(r, path=fd)132 133        print('{} saved to {}'.format(url, filename))134 135    .. code-block:: python136 137        import io138        import requests139        from requests_toolbelt.downloadutils import stream140 141        b = io.BytesIO()142        r = requests.get(url, stream=True)143        filename = stream.stream_response_to_file(r, path=b)144        assert filename is None145 146    :param response: A Response object from requests147    :type response: requests.models.Response148    :param path: *(optional)*, Either a string with the path to the location149        to save the response content, or a file-like object expecting bytes.150    :type path: :class:`str`, or object with a :meth:`write`151    :param int chunksize: (optional), Size of chunk to attempt to stream152        (default 512B).153    :returns: The name of the file, if one can be determined, else None154    :rtype: str155    :raises: :class:`requests_toolbelt.exceptions.StreamingError`156    """157    pre_opened = False158    fd = None159    filename = None160    if path and callable(getattr(path, 'write', None)):161        pre_opened = True162        fd = path163        filename = getattr(fd, 'name', None)164    else:165        filename = get_download_file_path(response, path)166        if os.path.exists(filename):167            raise exc.StreamingError("File already exists: %s" % filename)168        fd = open(filename, 'wb')169 170    for chunk in response.iter_content(chunk_size=chunksize):171        fd.write(chunk)172 173    if not pre_opened:174        fd.close()175 176    return filename177