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

by the urllib3 version
gdubicki · Jul 30, 2020 · 31caaba · 31caaba
1 parent 2d39c0d
commit 31caaba
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 1 deletion.
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,45 @@ 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
+        pybrotli = True
+    except ImportError:
+        pybrotli = False
+
+    return urllib3_with_brotli and pybrotli
+
+
+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,39 @@ def response_handler(sock):
         assert r.url == 'http://{}:{}/final-url/#relevant-section'.format(host, port)
 
         close_server.set()
+
+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.
+    """
+
+    if brotli_supported():
+
+        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()