Accept Brotli as encoding, if supported (psf#4525)

by the urllib3 version
gdubicki · Aug 9, 2020 · 329c775 · 329c775
1 parent 2d39c0d
commit 329c775
Show file tree

Hide file tree

Showing 6 changed files with 98 additions and 2 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -21,6 +21,11 @@ jobs:
       dist: xenial
     - stage: test
       python: '3.8'
+    - stage: test
+      python: '3.8'
+      name: '3.8 without Brotli'
+      before_script:
+        - pip uninstall brotli
     - stage: coverage
       python: '3.6'
       script: codecov
diff --git a/Pipfile b/Pipfile
@@ -22,4 +22,4 @@ sphinx = "<=1.5.5"
 tox = "*"
 
 [packages]
-"requests" = {path = ".", editable = true, extras = ["socks"]}
+"requests" = {path = ".", editable = true, extras = ["socks", "brotli"]}
diff --git a/requests/utils.py b/requests/utils.py
@@ -20,6 +20,7 @@
 import warnings
 import zipfile
 from collections import OrderedDict
+import urllib3
 
 from .__version__ import __version__
 from . import certs
@@ -804,13 +805,44 @@ def default_user_agent(name="python-requests"):
     return '%s/%s' % (name, __version__)
 
 
+def brotli_supported():
+    """
+    Returns whether Brotli compression is supported
+
+    :rtype: bool
+    """
+
+    # urllib >= 1.25.1 includes brotli support
+    major, minor, patch = urllib3.__version__.split('.')  # noqa: F811
+    major, minor, patch = int(major), int(minor), int(patch)
+    urllib3_with_brotli = (major == 1 and ((minor == 25 and patch >= 1) or (minor >= 26))) \
+        or major >= 2
+
+    # pybrotli is an extra package required by urllib3 for brotli support
+    try:
+        import brotli
+    except ImportError:
+        brotli = None
+
+    return urllib3_with_brotli and brotli
+
+
+BROTLI_SUPPORTED = brotli_supported()
+
+
 def default_headers():
     """
     :rtype: requests.structures.CaseInsensitiveDict
     """
+
+    if BROTLI_SUPPORTED:
+        accepted_encodings = ('br', 'gzip', 'deflate')
+    else:
+        accepted_encodings = ('gzip', 'deflate')
+
     return CaseInsensitiveDict({
         'User-Agent': default_user_agent(),
-        'Accept-Encoding': ', '.join(('gzip', 'deflate')),
+        'Accept-Encoding': ', '.join(accepted_encodings),
         'Accept': '*/*',
         'Connection': 'keep-alive',
     })

diff --git a/setup.py b/setup.py
@@ -105,6 +105,7 @@ def run_tests(self):
         'security': ['pyOpenSSL >= 0.14', 'cryptography>=1.3.4'],
         'socks': ['PySocks>=1.5.6, !=1.5.7'],
         'socks:sys_platform == "win32" and python_version == "2.7"': ['win_inet_pton'],
+        'brotli': ['urllib3 >= 1.25.1', 'urllib3[brotli]'],
     },
     project_urls={
         'Documentation': 'https://requests.readthedocs.io',

diff --git a/tests/test_lowlevel.py b/tests/test_lowlevel.py
@@ -3,7 +3,9 @@
 import pytest
 import threading
 import requests
+import brotli
 
+from requests.utils import brotli_supported
 from tests.testserver.server import Server, consume_socket_content
 
 from .utils import override_environ
@@ -307,3 +309,38 @@ def response_handler(sock):
         assert r.url == 'http://{}:{}/final-url/#relevant-section'.format(host, port)
 
         close_server.set()
+
+@pytest.mark.skipif(not brotli_supported(), reason='run only if Brotli is supported')
+def test_brotli():
+    """Verify that if Requests supports Brotli, then a request to a server serving
+    content in Brotli will result with a response with the content correctly decompressed.
+    """
+
+    some_text = '<html><html><html>'
+    br_compressed_text = brotli.compress(some_text.encode('utf-8'))
+
+    assert len(some_text) != len(br_compressed_text)
+    assert some_text != br_compressed_text
+
+    def response_handler(sock):
+        consume_socket_content(sock, timeout=0.5)
+
+        sock.send(
+            b'HTTP/1.1 200 OK\r\n'
+            b'Content-Length: ' + bytes(len(br_compressed_text)) + b'\r\n'
+            b'Content-Encoding: br\r\n'
+            b'\r\n' +
+            br_compressed_text
+        )
+
+    close_server = threading.Event()
+    server = Server(response_handler, wait_to_close_event=close_server)
+
+    with server as (host, port):
+        url = 'http://{}:{}/'.format(host, port)
+        r = requests.get(url)
+
+        assert len(r.text) == len(some_text)
+        assert r.text == some_text
+
+        close_server.set()
diff --git a/tests/test_requests.py b/tests/test_requests.py
@@ -31,6 +31,7 @@
 from requests.models import urlencode
 from requests.hooks import default_hooks
 from requests.compat import MutableMapping
+from requests.utils import brotli_supported
 
 from .compat import StringIO, u
 from .utils import override_environ
@@ -2527,3 +2528,23 @@ def test_parameters_for_nonstandard_schemes(self, input, params, expected):
         r = requests.Request('GET', url=input, params=params)
         p = r.prepare()
         assert p.url == expected
+
+
+    @pytest.mark.skipif(not brotli_supported(), reason='run only if Brotli is supported')
+    def test_brotli_default_accept_encoding(self, httpbin):
+        """Verify that Requests with Brotli support by default sends Accept-Encoding with 'br',
+        as the first accepted encoding as the most effecient
+        """
+        r = requests.request('GET', httpbin('get'))
+
+        assert 'br' in r.request.headers['Accept-Encoding']
+        assert r.request.headers['Accept-Encoding'].startswith('br')
+
+
+    @pytest.mark.skipif(brotli_supported(), reason='run only if Brotli is NOT supported')
+    def test_no_brotli_default_accept_encoding(self, httpbin):
+        """Verify that Requests without Brotli support by default sends Accept-Encoding without 'br'
+        """
+        r = requests.request('GET', httpbin('get'))
+
+        assert 'br' not in r.request.headers['Accept-Encoding']