Merge pull request #4831 from starrify/downloadermw-support-zstd

Adding support for zstd in HttpCompressionMiddleware

docs/topics/downloader-middleware.rst

@@ -684,11 +684,14 @@ HttpCompressionMiddleware
    This middleware allows compressed (gzip, deflate) traffic to be
    sent/received from web sites.
 
-   This middleware also supports decoding `brotli-compressed`_ responses,
-   provided `brotlipy`_ is installed.
+   This middleware also supports decoding `brotli-compressed`_ as well as
+   `zstd-compressed`_ responses, provided that `brotlipy`_ or `zstandard`_ is
+   installed, respectively.
 
 .. _brotli-compressed: https://www.ietf.org/rfc/rfc7932.txt
 .. _brotlipy: https://pypi.org/project/brotlipy/
+.. _zstd-compressed: https://www.ietf.org/rfc/rfc8478.txt
+.. _zstandard: https://pypi.org/project/zstandard/
 
 HttpCompressionMiddleware Settings
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

scrapy/downloadermiddlewares/httpcompression.py

@@ -1,3 +1,4 @@
+import io
 import zlib
 
 from scrapy.utils.gz import gunzip
@@ -14,6 +15,12 @@ try:
 except ImportError:
     pass
 
+try:
+    import zstandard
+    ACCEPTED_ENCODINGS.append(b'zstd')
+except ImportError:
+    pass
+
 
 class HttpCompressionMiddleware:
     """This middleware allows compressed (gzip, deflate) traffic to be
@@ -67,4 +74,9 @@ class HttpCompressionMiddleware:
                 body = zlib.decompress(body, -15)
         if encoding == b'br' and b'br' in ACCEPTED_ENCODINGS:
             body = brotli.decompress(body)
+        if encoding == b'zstd' and b'zstd' in ACCEPTED_ENCODINGS:
+            # Using its streaming API since its simple API could handle only cases
+            # where there is content size data embedded in the frame
+            reader = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(body))
+            body = reader.read()
         return body

tests/requirements-py3.txt

@@ -16,6 +16,7 @@ uvloop; platform_system != "Windows"
 
 # optional for shell wrapper tests
 bpython
-brotlipy
+brotlipy  # optional for HTTP compress downloader middleware tests
+zstandard  # optional for HTTP compress downloader middleware tests
 ipython
 pywin32; sys_platform == "win32"

tests/test_downloadermiddleware_httpcompression.py

@@ -20,6 +20,12 @@ FORMAT = {
     'rawdeflate': ('html-rawdeflate.bin', 'deflate'),
     'zlibdeflate': ('html-zlibdeflate.bin', 'deflate'),
     'br': ('html-br.bin', 'br'),
+    # $ zstd raw.html --content-size -o html-zstd-static-content-size.bin
+    'zstd-static-content-size': ('html-zstd-static-content-size.bin', 'zstd'),
+    # $ zstd raw.html --no-content-size -o html-zstd-static-no-content-size.bin
+    'zstd-static-no-content-size': ('html-zstd-static-no-content-size.bin', 'zstd'),
+    # $ cat raw.html | zstd -o html-zstd-streaming-no-content-size.bin
+    'zstd-streaming-no-content-size': ('html-zstd-streaming-no-content-size.bin', 'zstd'),
 }
 
 
@@ -80,6 +86,27 @@ class HttpCompressionTest(TestCase):
         assert newresponse.body.startswith(b"<!DOCTYPE")
         assert 'Content-Encoding' not in newresponse.headers
 
+    def test_process_response_zstd(self):
+        try:
+            import zstandard  # noqa: F401
+        except ImportError:
+            raise SkipTest("no zstd support (zstandard)")
+        raw_content = None
+        for check_key in FORMAT:
+            if not check_key.startswith('zstd-'):
+                continue
+            response = self._getresponse(check_key)
+            request = response.request
+            self.assertEqual(response.headers['Content-Encoding'], b'zstd')
+            newresponse = self.mw.process_response(request, response, self.spider)
+            if raw_content is None:
+                raw_content = newresponse.body
+            else:
+                assert raw_content == newresponse.body
+            assert newresponse is not response
+            assert newresponse.body.startswith(b"<!DOCTYPE")
+            assert 'Content-Encoding' not in newresponse.headers
+
     def test_process_response_rawdeflate(self):
         response = self._getresponse('rawdeflate')
         request = response.request