Add support for Zstandard compression

Author: kge

CHANGES/11161.feature.rst

@@ -0,0 +1,2 @@
+Add support for Zstandard (aka Zstd) compression
+-- by :user:`KGuillaume-chaps`.

CONTRIBUTORS.txt

@@ -218,6 +218,7 @@ Justin Foo
 Justin Turner Arthur
 Kay Zheng
 Kevin Samuel
+Kilian Guillaume
 Kimmo Parviainen-Jalanko
 Kirill Klenov
 Kirill Malovitsa

aiohttp/_http_parser.pyx

@@ -437,7 +437,7 @@ cdef class HttpParser:
         if enc is not None:
             self._content_encoding = None
             enc = enc.lower()
-            if enc in ('gzip', 'deflate', 'br'):
+            if enc in ('gzip', 'deflate', 'br', 'zstd'):
                 encoding = enc
 
         if self._cparser.type == cparser.HTTP_REQUEST:

aiohttp/client_reqrep.py

@@ -40,7 +40,7 @@
     InvalidURL,
     ServerFingerprintMismatch,
 )
-from .compression_utils import HAS_BROTLI
+from .compression_utils import HAS_BROTLI, HAS_ZSTD
 from .formdata import FormData
 from .hdrs import CONTENT_TYPE
 from .helpers import (
@@ -101,7 +101,15 @@
 
 
 def _gen_default_accept_encoding() -> str:
-    return "gzip, deflate, br" if HAS_BROTLI else "gzip, deflate"
+    encodings = [
+        "gzip",
+        "deflate",
+    ]
+    if HAS_BROTLI:
+        encodings.append("br")
+    if HAS_ZSTD:
+        encodings.append("zstd")
+    return ", ".join(encodings)
 
 
 @frozen_dataclass_decorator

aiohttp/compression_utils.py

@@ -21,6 +21,13 @@
 except ImportError:
     HAS_BROTLI = False
 
+try:
+    import zstandard
+
+    HAS_ZSTD = True
+except ImportError:
+    HAS_ZSTD = False
+
 MAX_SYNC_CHUNK_SIZE = 1024
 
 
@@ -276,3 +283,19 @@ def flush(self) -> bytes:
         if hasattr(self._obj, "flush"):
             return cast(bytes, self._obj.flush())
         return b""
+
+
+class ZSTDDecompressor:
+    def __init__(self) -> None:
+        if not HAS_ZSTD:
+            raise RuntimeError(
+                "The zstd decompression is not available. "
+                "Please install `zstandard` module"
+            )
+        self._obj = zstandard.ZstdDecompressor()
+
+    def decompress_sync(self, data: bytes) -> bytes:
+        return cast(bytes, self._obj.decompress(data))
+
+    def flush(self) -> bytes:
+        return b""

aiohttp/http_parser.py

@@ -26,7 +26,13 @@
 
 from . import hdrs
 from .base_protocol import BaseProtocol
-from .compression_utils import HAS_BROTLI, BrotliDecompressor, ZLibDecompressor
+from .compression_utils import (
+    HAS_BROTLI,
+    HAS_ZSTD,
+    BrotliDecompressor,
+    ZLibDecompressor,
+    ZSTDDecompressor,
+)
 from .helpers import (
     _EXC_SENTINEL,
     DEBUG,
@@ -527,7 +533,7 @@ def parse_headers(
         enc = headers.get(hdrs.CONTENT_ENCODING)
         if enc:
             enc = enc.lower()
-            if enc in ("gzip", "deflate", "br"):
+            if enc in ("gzip", "deflate", "br", "zstd"):
                 encoding = enc
 
         # chunking
@@ -930,14 +936,21 @@ def __init__(self, out: StreamReader, encoding: Optional[str]) -> None:
         self.encoding = encoding
         self._started_decoding = False
 
-        self.decompressor: Union[BrotliDecompressor, ZLibDecompressor]
+        self.decompressor: Union[BrotliDecompressor, ZLibDecompressor, ZSTDDecompressor]
         if encoding == "br":
             if not HAS_BROTLI:
                 raise ContentEncodingError(
                     "Can not decode content-encoding: brotli (br). "
                     "Please install `Brotli`"
                 )
             self.decompressor = BrotliDecompressor()
+        elif encoding == "zstd":
+            if not HAS_ZSTD:
+                raise ContentEncodingError(
+                    "Can not decode content-encoding: zstandard (zstd). "
+                    "Please install `zstandard`"
+                )
+            self.decompressor = ZSTDDecompressor()
         else:
             self.decompressor = ZLibDecompressor(encoding=encoding)
 

docs/client_quickstart.rst

@@ -190,6 +190,9 @@ You can enable ``brotli`` transfer-encodings support,
 just install `Brotli <https://pypi.org/project/Brotli/>`_
 or `brotlicffi <https://pypi.org/project/brotlicffi/>`_.
 
+You can enable ``zstd`` transfer-encodings support,
+install `zstandard <https://pypi.org/project/zstandard/>`_.
+
 JSON Request
 ============
 

docs/spelling_wordlist.txt

@@ -381,3 +381,5 @@ www
 xxx
 yarl
 zlib
+zstandard
+zstd

requirements/runtime-deps.in

@@ -10,3 +10,4 @@ frozenlist >= 1.1.1
 multidict >=4.5, < 7.0
 propcache >= 0.2.0
 yarl >= 1.17.0, < 2.0
+zstandard; platform_python_implementation == 'CPython'

setup.cfg

@@ -67,6 +67,7 @@ speedups =
   aiodns >= 3.3.0
   Brotli; platform_python_implementation == 'CPython'
   brotlicffi; platform_python_implementation != 'CPython'
+  zstandard; platform_python_implementation == 'CPython'
 
 [options.packages.find]
 exclude =

tests/test_client_request.py

@@ -358,7 +358,7 @@ def test_headers(make_request: _RequestMaker) -> None:
 
     assert hdrs.CONTENT_TYPE in req.headers
     assert req.headers[hdrs.CONTENT_TYPE] == "text/plain"
-    assert req.headers[hdrs.ACCEPT_ENCODING] == "gzip, deflate, br"
+    assert req.headers[hdrs.ACCEPT_ENCODING] == "gzip, deflate, br, zstd"
 
 
 def test_headers_list(make_request: _RequestMaker) -> None:
@@ -1568,15 +1568,20 @@ def test_loose_cookies_types(loop: asyncio.AbstractEventLoop) -> None:
 
 
 @pytest.mark.parametrize(
-    "has_brotli,expected",
+    "has_brotli,has_zstd,expected",
     [
-        (False, "gzip, deflate"),
-        (True, "gzip, deflate, br"),
+        (False, False, "gzip, deflate"),
+        (True, False, "gzip, deflate, br"),
+        (False, True, "gzip, deflate, zstd"),
+        (True, True, "gzip, deflate, br, zstd"),
     ],
 )
-def test_gen_default_accept_encoding(has_brotli: bool, expected: str) -> None:
+def test_gen_default_accept_encoding(
+    has_brotli: bool, has_zstd: bool, expected: str
+) -> None:
     with mock.patch("aiohttp.client_reqrep.HAS_BROTLI", has_brotli):
-        assert _gen_default_accept_encoding() == expected
+        with mock.patch("aiohttp.client_reqrep.HAS_ZSTD", has_zstd):
+            assert _gen_default_accept_encoding() == expected
 
 
 @pytest.mark.parametrize(

tests/test_http_parser.py

@@ -34,6 +34,11 @@
 except ImportError:
     brotli = None
 
+try:
+    import zstandard
+except ImportError:
+    zstandard = None
+
 REQUEST_PARSERS = [HttpRequestParserPy]
 RESPONSE_PARSERS = [HttpResponseParserPy]
 
@@ -600,6 +605,14 @@ def test_compression_brotli(parser: HttpRequestParser) -> None:
     assert msg.compression == "br"
 
 
+@pytest.mark.skipif(zstandard is None, reason="zstandard is not installed")
+def test_compression_zstd(parser: HttpRequestParser) -> None:
+    text = b"GET /test HTTP/1.1\r\ncontent-encoding: zstd\r\n\r\n"
+    messages, upgrade, tail = parser.feed_data(text)
+    msg = messages[0][0]
+    assert msg.compression == "zstd"
+
+
 def test_compression_unknown(parser: HttpRequestParser) -> None:
     text = b"GET /test HTTP/1.1\r\ncontent-encoding: compress\r\n\r\n"
     messages, upgrade, tail = parser.feed_data(text)
@@ -1849,6 +1862,15 @@ async def test_http_payload_brotli(self, protocol: BaseProtocol) -> None:
         assert b"brotli data" == out._buffer[0]
         assert out.is_eof()
 
+    @pytest.mark.skipif(zstandard is None, reason="zstandard is not installed")
+    async def test_http_payload_zstandard(self, protocol: BaseProtocol) -> None:
+        compressed = zstandard.compress(b"zstd data")
+        out = aiohttp.StreamReader(protocol, 2**16, loop=asyncio.get_running_loop())
+        p = HttpPayloadParser(out, length=len(compressed), compression="zstd")
+        p.feed_data(compressed)
+        assert b"zstd data" == out._buffer[0]
+        assert out.is_eof()
+
 
 class TestDeflateBuffer:
     async def test_feed_data(self, protocol: BaseProtocol) -> None:
@@ -1919,6 +1941,17 @@ async def test_feed_eof_no_err_brotli(self, protocol: BaseProtocol) -> None:
         dbuf.feed_eof()
         assert [b"line"] == list(buf._buffer)
 
+    async def test_feed_eof_no_err_zstandard(self, protocol: BaseProtocol) -> None:
+        buf = aiohttp.StreamReader(protocol, 2**16, loop=asyncio.get_running_loop())
+        dbuf = DeflateBuffer(buf, "zstd")
+
+        dbuf.decompressor = mock.Mock()
+        dbuf.decompressor.flush.return_value = b"line"
+        dbuf.decompressor.eof = False
+
+        dbuf.feed_eof()
+        assert [b"line"] == list(buf._buffer)
+
     async def test_empty_body(self, protocol: BaseProtocol) -> None:
         buf = aiohttp.StreamReader(protocol, 2**16, loop=asyncio.get_running_loop())
         dbuf = DeflateBuffer(buf, "deflate")