|
6 | 6 |
|
7 | 7 | from .util import CaseInsensitiveDict
|
8 | 8 |
|
| 9 | +try: |
| 10 | + # This supports both brotli & brotlipy packages |
| 11 | + import brotli |
| 12 | +except ImportError: |
| 13 | + try: |
| 14 | + import brotlicffi as brotli |
| 15 | + except ImportError: |
| 16 | + brotli = None |
| 17 | + |
| 18 | + |
| 19 | +AVAILABLE_DECOMPRESSORS = {"gzip", "deflate"} |
| 20 | +if brotli is not None: |
| 21 | + AVAILABLE_DECOMPRESSORS.add("br") |
| 22 | + |
9 | 23 |
|
10 | 24 | def replace_headers(request, replacements):
|
11 | 25 | """Replace headers in request according to replacements.
|
@@ -136,30 +150,33 @@ def remove_post_data_parameters(request, post_data_parameters_to_remove):
|
136 | 150 |
|
137 | 151 | def decode_response(response):
|
138 | 152 | """
|
139 |
| - If the response is compressed with gzip or deflate: |
| 153 | + If the response is compressed with any supported compression (gzip, |
| 154 | + deflate, br if available): |
140 | 155 | 1. decompress the response body
|
141 | 156 | 2. delete the content-encoding header
|
142 | 157 | 3. update content-length header to decompressed length
|
143 | 158 | """
|
144 | 159 |
|
145 |
| - def is_compressed(headers): |
| 160 | + def is_decompressable(headers): |
146 | 161 | encoding = headers.get("content-encoding", [])
|
147 |
| - return encoding and encoding[0] in ("gzip", "deflate") |
| 162 | + return encoding and encoding[0] in AVAILABLE_DECOMPRESSORS |
148 | 163 |
|
149 | 164 | def decompress_body(body, encoding):
|
150 | 165 | """Returns decompressed body according to encoding using zlib.
|
151 | 166 | to (de-)compress gzip format, use wbits = zlib.MAX_WBITS | 16
|
152 | 167 | """
|
153 | 168 | if encoding == "gzip":
|
154 | 169 | return zlib.decompress(body, zlib.MAX_WBITS | 16)
|
155 |
| - else: # encoding == 'deflate' |
| 170 | + elif encoding == "deflate": |
156 | 171 | return zlib.decompress(body)
|
| 172 | + else: # encoding == 'br' |
| 173 | + return brotli.decompress(body) |
157 | 174 |
|
158 | 175 | # Deepcopy here in case `headers` contain objects that could
|
159 | 176 | # be mutated by a shallow copy and corrupt the real response.
|
160 | 177 | response = copy.deepcopy(response)
|
161 | 178 | headers = CaseInsensitiveDict(response["headers"])
|
162 |
| - if is_compressed(headers): |
| 179 | + if is_decompressable(headers): |
163 | 180 | encoding = headers["content-encoding"][0]
|
164 | 181 | headers["content-encoding"].remove(encoding)
|
165 | 182 | if not headers["content-encoding"]:
|
|
0 commit comments