feat: zstd middleware (#4286)

Author: ceballots

litestar/config/compression.py

@@ -21,7 +21,7 @@ class CompressionConfig:
     using the ``compression_config`` key.
     """
 
-    backend: Literal["gzip", "brotli"] | str
+    backend: Literal["gzip", "brotli", "zstd"] | str
     """The backend to use.
 
     If the value given is `gzip` or `brotli`, then the builtin gzip and brotli compression is used.
@@ -30,6 +30,8 @@ class CompressionConfig:
     """Minimum response size (bytes) to enable compression, affects all backends."""
     gzip_compress_level: int = field(default=9)
     """Range ``[0-9]``, see :doc:`python:library/gzip`."""
+    zstd_compress_level: int = field(default=3)
+    """Range `[0-22]`, see `zstandard <https://pypi.org/project/zstandard/>`"""
     brotli_quality: int = field(default=5)
     """Range ``[0-11]``, Controls the compression-speed vs compression-density tradeoff.
 
@@ -59,6 +61,8 @@ class CompressionConfig:
     """The compression facade to use for the actual compression."""
     backend_config: Any = None
     """Configuration specific to the backend."""
+    zstd_gzip_fallback: bool = True
+    """Use GZIP as a fallback if Zstd is not supported by the client."""
     gzip_fallback: bool = True
     """Use GZIP as a fallback if the provided backend is not supported by the client."""
 
@@ -81,3 +85,13 @@ def __post_init__(self) -> None:
 
             self.gzip_fallback = self.brotli_gzip_fallback
             self.compression_facade = BrotliCompression
+        elif self.backend == "zstd":
+            if self.zstd_compress_level < 1 or self.zstd_compress_level > 22:
+                raise ImproperlyConfiguredException(
+                    f"zstd_compress_level must be between 1 and 22, given: {self.zstd_compress_level}"
+                )
+
+            from litestar.middleware.compression.zstd_facade import ZstdCompression
+
+            self.gzip_fallback = self.zstd_gzip_fallback
+            self.compression_facade = ZstdCompression

litestar/enums.py

@@ -77,6 +77,7 @@ class CompressionEncoding(str, Enum):
 
     GZIP = "gzip"
     BROTLI = "br"
+    ZSTD = "zstd"
 
 
 class ASGIExtension(str, Enum):

litestar/middleware/compression/brotli_facade.py

@@ -43,7 +43,7 @@ def __init__(
             lgblock=config.brotli_lgblock,
         )
 
-    def write(self, body: bytes) -> None:
+    def write(self, body: bytes | bytearray, final: bool = False) -> None:
         self.buffer.write(self.compressor.process(body))
         self.buffer.write(self.compressor.flush())
 

litestar/middleware/compression/facade.py

@@ -29,11 +29,13 @@ def __init__(
         """
         ...
 
-    def write(self, body: bytes) -> None:
-        """Write compressed bytes.
+    def write(self, body: bytes | bytearray, final: bool = False) -> None:
+        """Write compressed bytes to the buffer.
 
         Args:
-            body: Message body to process
+            body: The message body to process. Can be `bytes` or `bytearray`.
+            final: Indicates whether this is the last chunk of data. If True,
+                the compressor may flush any remaining internal buffers.
 
         Returns:
             None

litestar/middleware/compression/gzip_facade.py

@@ -24,8 +24,9 @@ def __init__(
         self.compression_encoding = compression_encoding
         self.compressor = GzipFile(mode="wb", fileobj=buffer, compresslevel=config.gzip_compress_level)
 
-    def write(self, body: bytes) -> None:
-        self.compressor.write(body)
+    def write(self, body: bytes | bytearray, final: bool = False) -> None:
+        data = bytes(body)
+        self.compressor.write(data)
         self.compressor.flush()
 
     def close(self) -> None:

litestar/middleware/compression/middleware.py

@@ -7,6 +7,7 @@
 from litestar.enums import CompressionEncoding, ScopeType
 from litestar.middleware.base import AbstractMiddleware
 from litestar.middleware.compression.gzip_facade import GzipCompression
+from litestar.middleware.compression.zstd_facade import ZstdCompression
 from litestar.utils.empty import value_or_default
 from litestar.utils.scope.state import ScopeState
 
@@ -82,10 +83,22 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
 
         await self.app(scope, receive, send)
 
+    def get_facade_cls(
+        self,
+        compression_encoding: Literal[CompressionEncoding.BROTLI, CompressionEncoding.GZIP, CompressionEncoding.ZSTD]
+        | str,
+    ) -> type[CompressionFacade]:
+        if compression_encoding == CompressionEncoding.GZIP:
+            return GzipCompression
+        if compression_encoding == CompressionEncoding.ZSTD:
+            return ZstdCompression
+        return self.config.compression_facade
+
     def create_compression_send_wrapper(
         self,
         send: Send,
-        compression_encoding: Literal[CompressionEncoding.BROTLI, CompressionEncoding.GZIP] | str,
+        compression_encoding: Literal[CompressionEncoding.BROTLI, CompressionEncoding.GZIP, CompressionEncoding.ZSTD]
+        | str,
         scope: Scope,
     ) -> Send:
         """Wrap ``send`` to handle brotli compression.
@@ -100,15 +113,10 @@ def create_compression_send_wrapper(
         """
         bytes_buffer = BytesIO()
 
-        facade: CompressionFacade
         # We can't use `self.config.compression_facade` directly if the compression is `gzip` since
         # it may be being used as a fallback.
-        if compression_encoding == CompressionEncoding.GZIP:
-            facade = GzipCompression(buffer=bytes_buffer, compression_encoding=compression_encoding, config=self.config)
-        else:
-            facade = self.config.compression_facade(
-                buffer=bytes_buffer, compression_encoding=compression_encoding, config=self.config
-            )
+        facade_cls: type[CompressionFacade] = self.get_facade_cls(compression_encoding)
+        facade = facade_cls(buffer=bytes_buffer, compression_encoding=compression_encoding, config=self.config)
 
         initial_message: HTTPResponseStartEvent | None = None
         started = False
@@ -151,7 +159,7 @@ async def send_wrapper(message: Message) -> None:
                         del headers["Content-Length"]
                         connection_state.response_compressed = True
 
-                        facade.write(body)
+                        facade.write(body, final=not more_body)
 
                         message["body"] = bytes_buffer.getvalue()
                         bytes_buffer.seek(0)
@@ -160,7 +168,7 @@ async def send_wrapper(message: Message) -> None:
                         await send(message)
 
                     elif len(body) >= self.config.minimum_size:
-                        facade.write(body)
+                        facade.write(body, final=not more_body)
                         facade.close()
                         body = bytes_buffer.getvalue()
 
@@ -180,7 +188,7 @@ async def send_wrapper(message: Message) -> None:
                         await send(message)
 
                 else:
-                    facade.write(body)
+                    facade.write(body, final=not more_body)
                     if not more_body:
                         facade.close()
 

litestar/middleware/compression/zstd_facade.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from litestar.enums import CompressionEncoding
+from litestar.exceptions import MissingDependencyException
+from litestar.middleware.compression.facade import CompressionFacade
+
+try:
+    import zstandard as zstd
+except ImportError as e:
+    raise MissingDependencyException("zstandard") from e
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from litestar.config.compression import CompressionConfig
+
+
+class ZstdCompression(CompressionFacade):
+    __slots__ = ("buffer", "cctx", "compression_encoding", "compressor")
+
+    encoding = CompressionEncoding("zstd")
+
+    def __init__(self, buffer: BytesIO, compression_encoding: Literal["zstd"] | str, config: CompressionConfig) -> None:
+        self.buffer = buffer
+        self.compression_encoding = compression_encoding
+        self.cctx = zstd.ZstdCompressor(level=config.zstd_compress_level)
+        self.compressor = self.cctx.stream_writer(buffer)
+
+    def write(self, body: bytes | bytearray, final: bool = False) -> None:
+        self.compressor.write(body)
+        if final:
+            self.compressor.flush(zstd.FLUSH_FRAME)
+        else:
+            self.compressor.flush(zstd.FLUSH_BLOCK)
+
+    def close(self) -> None:
+        self.compressor.flush(zstd.FLUSH_FRAME)

pyproject.toml

@@ -56,6 +56,10 @@ dependencies = [
 annotated-types = ["annotated-types"]
 attrs = ["attrs"]
 brotli = ["brotli"]
+zstd = [
+    # Only needed for <3.14, since >=3.14 ships with stdlib support
+    "zstandard>=0.24.0; python_version < '3.14'",
+]
 cli = ["jsbeautifier", "uvicorn[standard]", "uvloop>=0.18.0; sys_platform != 'win32'"]
 cryptography = ["cryptography"]
 full = [
@@ -80,7 +84,8 @@ full = [
     structlog, \
     valkey, \
     htmx, \
-    yaml]; \
+    yaml, \
+    zstd]; \
     python_version < \"3.13\"""",
 
   # full extra without picologging
@@ -104,7 +109,8 @@ full = [
     structlog, \
     valkey, \
     htmx, \
-    yaml]; \
+    yaml, \
+    zstd]; \
     python_version >= \"3.13\"""",
 ]
 htmx = ["litestar-htmx>=0.4.0"]
@@ -329,6 +335,8 @@ module = [
   "exceptiongroup",
   "picologging",
   "picologging.*",
+  "zstandard",
+  "zstandard.*"
 ]
 
 [[tool.mypy.overrides]]

tests/e2e/test_logging/test_structlog_to_file.py

@@ -61,7 +61,7 @@ def handler() -> str:
             "headers": {
                 "host": "testserver.local",
                 "accept": "*/*",
-                "accept-encoding": "gzip, deflate, br",
+                "accept-encoding": "gzip, deflate, br, zstd",
                 "connection": "keep-alive",
                 "user-agent": "testclient",
             },

tests/unit/test_connection/test_request.py

@@ -171,7 +171,7 @@ async def app(scope: Scope, receive: Receive, send: Send) -> None:
         "headers": {
             "host": "example.org",
             "user-agent": "testclient",
-            "accept-encoding": "gzip, deflate, br",
+            "accept-encoding": "gzip, deflate, br, zstd",
             "accept": "*/*",
             "connection": "keep-alive",
         }