Merge pull request #3897 from nateprewitt/prep_body

Refactor Prepare Body

3.0-HISTORY.rst

@@ -1,6 +1,10 @@
 3.0.0 (2017-xx-xx)
 ++++++++++++++++++
 
+- Simplified logic for determining Content-Length and Transfer-Encoding.
+  Requests will now avoid setting both headers on the same request, and
+  raise an exception if this is done manually by a user.
+
 - Remove the HTTPProxyAuth class in favor of supporting proxy auth via
   the proxies parameter.
 

requests/exceptions.py

@@ -104,9 +104,14 @@ class StreamConsumedError(RequestException, TypeError):
 class RetryError(RequestException):
     """Custom retries logic failed"""
 
+
 class UnrewindableBodyError(RequestException):
     """Requests encountered an error when trying to rewind a body"""
 
+
+class InvalidBodyError(RequestException, ValueError):
+    """An invalid request body was specified"""
+
 # Warnings
 
 

requests/models.py

@@ -32,12 +32,14 @@ from .packages.urllib3.exceptions import (
     LocationParseError, ConnectionError)
 from .exceptions import (
     HTTPError, MissingScheme, InvalidURL, ChunkedEncodingError,
-    ContentDecodingError, ConnectionError, StreamConsumedError)
+    ContentDecodingError, ConnectionError, StreamConsumedError,
+    InvalidHeader, InvalidBodyError)
 from ._internal_utils import to_native_string, unicode_is_ascii
 from .utils import (
     guess_filename, get_auth_from_url, requote_uri,
     stream_decode_response_unicode, to_key_val_list, parse_header_links,
-    iter_slices, guess_json_utf, super_len, check_header_validity)
+    iter_slices, guess_json_utf, super_len, check_header_validity,
+    is_stream)
 from .compat import (
     cookielib, urlunparse, urlsplit, urlencode, str, bytes, StringIO,
     is_py2, chardet, builtin_str, basestring)
@@ -466,17 +468,7 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin):
             if not isinstance(body, bytes):
                 body = body.encode('utf-8')
 
-        is_stream = all([
-            hasattr(data, '__iter__'),
-            not isinstance(data, (basestring, list, tuple, collections.Mapping))
-        ])
-
-        try:
-            length = super_len(data)
-        except (TypeError, AttributeError, UnsupportedOperation):
-            length = None
-
-        if is_stream:
+        if is_stream(data):
             body = data
 
             if getattr(body, 'tell', None) is not None:
@@ -493,10 +485,6 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin):
             if files:
                 raise NotImplementedError('Streamed bodies and files are mutually exclusive.')
 
-            if length:
-                self.headers['Content-Length'] = builtin_str(length)
-            else:
-                self.headers['Transfer-Encoding'] = 'chunked'
         else:
             # Multi-part file uploads.
             if files:
@@ -509,27 +497,40 @@ class PreparedRequest(RequestEncodingMixin, RequestHooksMixin):
                     else:
                         content_type = 'application/x-www-form-urlencoded'
 
-            self.prepare_content_length(body)
-
             # Add content-type if it wasn't explicitly provided.
             if content_type and ('content-type' not in self.headers):
                 self.headers['Content-Type'] = content_type
 
+        self.prepare_content_length(body)
         self.body = body
 
     def prepare_content_length(self, body):
-        """Prepare Content-Length header based on request method and body"""
+        """Prepares Content-Length header.
+        
+        If the length of the body of the request can be computed, Content-Length
+        is set using ``super_len``. If user has manually set either a
+        Transfer-Encoding or Content-Length header when it should not be set
+        (they should be mutually exclusive) an InvalidHeader
+        error will be raised.
+        """
         if body is not None:
             length = super_len(body)
+
             if length:
-                # If length exists, set it. Otherwise, we fallback
-                # to Transfer-Encoding: chunked.
                 self.headers['Content-Length'] = builtin_str(length)
+            elif is_stream(body):
+                self.headers['Transfer-Encoding'] = 'chunked'
+            else:
+                raise InvalidBodyError('Non-null body must have length or be streamable.')
         elif self.method not in ('GET', 'HEAD') and self.headers.get('Content-Length') is None:
             # Set Content-Length to 0 for methods that can have a body
             # but don't provide one. (i.e. not GET or HEAD)
             self.headers['Content-Length'] = '0'
 
+        if 'Transfer-Encoding' in self.headers and 'Content-Length' in self.headers:
+            raise InvalidHeader('Conflicting Headers: Both Transfer-Encoding and '
+                                'Content-Length are set.')
+
     def prepare_auth(self, auth, url=''):
         """Prepares the given HTTP auth data."""
 

requests/utils.py

@@ -865,6 +865,7 @@ def urldefragauth(url):
 
     return urlunparse((scheme, netloc, path, params, query, ''))
 
+
 def rewind_body(prepared_request):
     """Move file pointer back to its recorded starting position
     so it can be read again on redirect.
@@ -878,3 +879,10 @@ def rewind_body(prepared_request):
                                         "body for redirect.")
     else:
         raise UnrewindableBodyError("Unable to rewind request body for redirect.")
+
+
+def is_stream(data):
+    """Given data, determines if it should be sent as a stream."""
+    is_iterable = getattr(data, '__iter__', False)
+    is_io_type = not isinstance(data, (basestring, list, tuple, collections.Mapping))
+    return is_iterable and is_io_type

tests/test_requests.py

@@ -25,7 +25,7 @@ from requests.cookies import (
 from requests.exceptions import (
     ConnectionError, ConnectTimeout, InvalidScheme, InvalidURL,
     MissingScheme, ReadTimeout, Timeout, RetryError, TooManyRedirects,
-    ProxyError, InvalidHeader, UnrewindableBodyError)
+    ProxyError, InvalidHeader, UnrewindableBodyError, InvalidBodyError)
 from requests.models import PreparedRequest
 from requests.structures import CaseInsensitiveDict
 from requests.sessions import SessionRedirectMixin
@@ -1924,6 +1924,61 @@ class TestRequests:
         assert 'Transfer-Encoding' in prepared_request.headers
         assert 'Content-Length' not in prepared_request.headers
 
+    def test_chunked_upload_with_manually_set_content_length_header_raises_error(self, httpbin):
+        """Ensure that if a user manually sets a content length header, when
+        the data is chunked, that an InvalidHeader error is raised.
+        """
+        data = (i for i in [b'a', b'b', b'c']) 
+        url = httpbin('post')
+        with pytest.raises(InvalidHeader):
+            r = requests.post(url, data=data, headers={'Content-Length': 'foo'})
+
+    def test_content_length_with_manually_set_transfer_encoding_raises_error(self, httpbin):
+        """Ensure that if a user manually sets a Transfer-Encoding header when
+        data is not chunked that an InvalidHeader error is raised.
+        """
+        data = 'test data'
+        url = httpbin('post')
+        with pytest.raises(InvalidHeader):
+            r = requests.post(url, data=data, headers={'Transfer-Encoding': 'chunked'})
+
+    def test_null_body_does_not_raise_error(self, httpbin):
+        url = httpbin('post')
+        try:
+            requests.post(url, data=None)
+        except InvalidHeader:
+            pytest.fail('InvalidHeader error raised unexpectedly.')
+
+    @pytest.mark.parametrize(
+        'body, expected', (
+            (None, ('Content-Length', '0')),
+            ('test_data', ('Content-Length', '9')),
+            (io.BytesIO(b'test_data'), ('Content-Length', '9')),
+            (StringIO.StringIO(''), ('Transfer-Encoding', 'chunked'))
+        ))
+    def test_prepare_content_length(self, httpbin, body, expected):
+        """Test prepare_content_length creates expected header."""
+        prep = requests.PreparedRequest()
+        prep.headers = {}
+        prep.method = 'POST'
+
+        # Ensure Content-Length is set appropriately.
+        key, value = expected
+        prep.prepare_content_length(body)
+        assert prep.headers[key] == value
+
+    def test_prepare_content_length_with_bad_body(self, httpbin):
+        """Test prepare_content_length raises exception with unsendable body."""
+        # Initialize minimum required PreparedRequest.
+        prep = requests.PreparedRequest()
+        prep.headers = {}
+        prep.method = 'POST'
+
+        with pytest.raises(InvalidBodyError) as e:
+            # Send object that isn't iterable and has no accessible content.
+            prep.prepare_content_length(object())
+            assert "Non-null body must have length or be streamable." in str(e)
+
     def test_custom_redirect_mixin(self, httpbin):
         """Tests a custom mixin to overwrite ``get_redirect_target``.