codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2"""3 4requests_toolbelt.streaming_iterator5====================================6 7This holds the implementation details for the :class:`StreamingIterator`. It8is designed for the case where you, the user, know the size of the upload but9need to provide the data as an iterator. This class will allow you to specify10the size and stream the data without using a chunked transfer-encoding.11 12"""13from requests.utils import super_len14 15from .multipart.encoder import CustomBytesIO, encode_with16 17 18class StreamingIterator(object):19 20    """21    This class provides a way of allowing iterators with a known size to be22    streamed instead of chunked.23 24    In requests, if you pass in an iterator it assumes you want to use25    chunked transfer-encoding to upload the data, which not all servers26    support well. Additionally, you may want to set the content-length27    yourself to avoid this but that will not work. The only way to preempt28    requests using a chunked transfer-encoding and forcing it to stream the29    uploads is to mimic a very specific interace. Instead of having to know30    these details you can instead just use this class. You simply provide the31    size and iterator and pass the instance of StreamingIterator to requests32    via the data parameter like so:33 34    .. code-block:: python35 36        from requests_toolbelt import StreamingIterator37 38        import requests39 40        # Let iterator be some generator that you already have and size be41        # the size of the data produced by the iterator42 43        r = requests.post(url, data=StreamingIterator(size, iterator))44 45    You can also pass file-like objects to :py:class:`StreamingIterator` in46    case requests can't determize the filesize itself. This is the case with47    streaming file objects like ``stdin`` or any sockets. Wrapping e.g. files48    that are on disk with ``StreamingIterator`` is unnecessary, because49    requests can determine the filesize itself.50 51    Naturally, you should also set the `Content-Type` of your upload52    appropriately because the toolbelt will not attempt to guess that for you.53    """54 55    def __init__(self, size, iterator, encoding='utf-8'):56        #: The expected size of the upload57        self.size = int(size)58 59        if self.size < 0:60            raise ValueError(61                'The size of the upload must be a positive integer'62                )63 64        #: Attribute that requests will check to determine the length of the65        #: body. See bug #80 for more details66        self.len = self.size67 68        #: Encoding the input data is using69        self.encoding = encoding70 71        #: The iterator used to generate the upload data72        self.iterator = iterator73 74        if hasattr(iterator, 'read'):75            self._file = iterator76        else:77            self._file = _IteratorAsBinaryFile(iterator, encoding)78 79    def read(self, size=-1):80        return encode_with(self._file.read(size), self.encoding)81 82 83class _IteratorAsBinaryFile(object):84    def __init__(self, iterator, encoding='utf-8'):85        #: The iterator used to generate the upload data86        self.iterator = iterator87 88        #: Encoding the iterator is using89        self.encoding = encoding90 91        # The buffer we use to provide the correct number of bytes requested92        # during a read93        self._buffer = CustomBytesIO()94 95    def _get_bytes(self):96        try:97            return encode_with(next(self.iterator), self.encoding)98        except StopIteration:99            return b''100 101    def _load_bytes(self, size):102        self._buffer.smart_truncate()103        amount_to_load = size - super_len(self._buffer)104        bytes_to_append = True105 106        while amount_to_load > 0 and bytes_to_append:107            bytes_to_append = self._get_bytes()108            amount_to_load -= self._buffer.append(bytes_to_append)109 110    def read(self, size=-1):111        size = int(size)112        if size == -1:113            return b''.join(self.iterator)114 115        self._load_bytes(size)116        return self._buffer.read(size)117