asgi-compression 0.1.0__py3-none-any.whl

asgi_compression/__init__.py

@@ -0,0 +1,17 @@
+from .base import CompressionAlgorithm, ContentEncoding
+from .brotli import BrotliAlgorithm, BrotliMode
+from .gzip import GzipAlgorithm
+from .identity import IdentityAlgorithm
+from .middleware import CompressionMiddleware
+from .zstd import ZstdAlgorithm
+
+__all__ = [
+    "CompressionMiddleware",
+    "CompressionAlgorithm",
+    "ContentEncoding",
+    "GzipAlgorithm",
+    "BrotliAlgorithm",
+    "BrotliMode",
+    "IdentityAlgorithm",
+    "ZstdAlgorithm",
+]

asgi_compression/base.py

@@ -0,0 +1,140 @@
+import typing
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+
+from .types import ASGIApp, Headers, Message, Receive, Scope, Send
+
+DEFAULT_EXCLUDED_CONTENT_TYPES = ("text/event-stream",)
+DEFAULT_MINIMUM_SIZE = 500
+
+
+class ContentEncoding(str, Enum):
+    GZIP = "gzip"
+    BROTLI = "br"
+    ZSTD = "zstd"
+    IDENTITY = "identity"
+
+
+async def unattached_send(message: Message) -> typing.NoReturn:
+    raise RuntimeError("send awaitable not set")  # pragma: no cover
+
+
+class CompressionResponder(ABC):
+    """Base class for all compression responders."""
+
+    content_encoding: ContentEncoding
+
+    def __init__(self, app: ASGIApp, minimum_size: int) -> None:
+        self.app = app
+        self.minimum_size = minimum_size
+        self._send: Send = unattached_send
+        self._initial_message: Message = {}
+        self._started = False
+        self._content_encoding_set = False
+        self._content_type_is_excluded = False
+
+    async def __call__(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+    ) -> None:
+        self._send = send
+        await self.app(scope, receive, self.send_with_compression)
+
+    async def send_with_compression(self, message: Message) -> None:
+        message_type = message["type"]
+        if message_type == "http.response.start":
+            # Don't send the initial message until we've determined how to
+            # modify the outgoing headers correctly.
+            self._initial_message = message
+            headers = Headers(raw=self._initial_message["headers"])
+
+            self._content_encoding_set = "content-encoding" in headers
+            self._content_type_is_excluded = headers.get(
+                "content-type", ""
+            ).startswith(DEFAULT_EXCLUDED_CONTENT_TYPES)
+
+        elif message_type == "http.response.body" and (
+            self._content_encoding_set or self._content_type_is_excluded
+        ):
+            if not self._started:
+                self._started = True
+                await self._send(self._initial_message)
+            await self._send(message)
+
+        elif message_type == "http.response.body" and not self._started:
+            self._started = True
+            body = message.get("body", b"")
+            more_body = message.get("more_body", False)
+
+            if len(body) < self.minimum_size and not more_body:
+                # Don't apply compression to small outgoing responses.
+                # Don't add Vary header for small responses
+                await self._send(self._initial_message)
+                await self._send(message)
+            elif not more_body:
+                # Standard response.
+                body = self.apply_compression(body, more_body=False)
+
+                headers = Headers(raw=self._initial_message["headers"])
+                headers.add_vary_header("Accept-Encoding")
+
+                if body != message["body"]:
+                    headers["Content-Encoding"] = self.content_encoding
+                    headers["Content-Length"] = str(len(body))
+                    message["body"] = body
+
+                self._initial_message["headers"] = headers.encode()
+                await self._send(self._initial_message)
+                await self._send(message)
+            else:
+                # Initial body in streaming response.
+                body = self.apply_compression(body, more_body=True)
+
+                headers = Headers(raw=self._initial_message["headers"])
+                headers.add_vary_header("Accept-Encoding")
+
+                if body != message["body"]:
+                    headers["Content-Encoding"] = self.content_encoding
+                    if "Content-Length" in headers:
+                        del headers["Content-Length"]
+
+                    message["body"] = body
+
+                self._initial_message["headers"] = headers.encode()
+                await self._send(self._initial_message)
+                await self._send(message)
+        elif message_type == "http.response.body":  # pragma: no branch
+            # Remaining body in streaming response.
+            body = message.get("body", b"")
+            more_body = message.get("more_body", False)
+
+            message["body"] = self.apply_compression(body, more_body=more_body)
+            await self._send(message)
+
+    @abstractmethod
+    def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
+        """Apply compression on the response body.
+
+        If more_body is False, any compression file should be closed. If it
+        isn't, it won't be closed automatically until all background tasks
+        complete.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class CompressionAlgorithm(ABC):
+    """Base class for compression algorithms."""
+
+    type: ContentEncoding
+    minimum_size: int = DEFAULT_MINIMUM_SIZE
+
+    def create_responder(self, app: ASGIApp) -> "CompressionResponder":
+        """Create a responder for this compression algorithm."""
+        raise NotImplementedError
+
+    def check_available(self) -> None:
+        """Check if the algorithm is available in the current environment."""

asgi_compression/brotli.py

@@ -0,0 +1,101 @@
+import io
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from .base import CompressionAlgorithm, CompressionResponder, ContentEncoding
+from .types import ASGIApp
+
+if TYPE_CHECKING:
+    import brotli
+
+
+def import_brotli() -> None:
+    global brotli
+    try:
+        import brotli
+    except ImportError as e:
+        raise ImportError(
+            "brotli is not installed, run `pip install brotli`]"
+        ) from e
+
+
+class BrotliMode(Enum):
+    TEXT = "text"
+    FONT = "font"
+    GENERIC = "generic"
+
+    def to_brotli_mode(self) -> int:
+        if self == BrotliMode.TEXT:
+            return brotli.MODE_TEXT
+        elif self == BrotliMode.FONT:
+            return brotli.MODE_FONT
+        elif self == BrotliMode.GENERIC:
+            return brotli.MODE_GENERIC
+        else:
+            assert False, f"Expected code to be unreachable, but got: {self}"
+
+
+class BrotliResponder(CompressionResponder):
+    """Responder that applies brotli compression."""
+
+    content_encoding = ContentEncoding.BROTLI
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        minimum_size: int,
+        quality: int = 4,
+        mode: BrotliMode = BrotliMode.TEXT,
+        lgwin: int = 22,
+        lgblock: int = 0,
+    ) -> None:
+        super().__init__(app, minimum_size)
+
+        import_brotli()
+
+        self.brotli_buffer = io.BytesIO()
+        self.compressor = brotli.Compressor(
+            quality=quality,
+            mode=mode.to_brotli_mode(),
+            lgwin=lgwin,
+            lgblock=lgblock,
+        )
+
+    def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
+        compressed = self.compressor.process(body)
+        self.brotli_buffer.write(compressed)
+
+        if not more_body:
+            final_data = self.compressor.finish()
+            self.brotli_buffer.write(final_data)
+
+        compressed_data = self.brotli_buffer.getvalue()
+
+        self.brotli_buffer.seek(0)
+        self.brotli_buffer.truncate()
+        return compressed_data
+
+
+@dataclass
+class BrotliAlgorithm(CompressionAlgorithm):
+    """Brotli compression algorithm."""
+
+    type: ContentEncoding = ContentEncoding.BROTLI
+    quality: int = 4
+    mode: BrotliMode = BrotliMode.TEXT
+    lgwin: int = 22
+    lgblock: int = 0
+
+    def create_responder(self, app: ASGIApp) -> "BrotliResponder":
+        return BrotliResponder(
+            app=app,
+            minimum_size=self.minimum_size,
+            quality=self.quality,
+            mode=self.mode,
+            lgwin=self.lgwin,
+            lgblock=self.lgblock,
+        )
+
+    def check_available(self) -> None:
+        import_brotli()

asgi_compression/gzip.py

@@ -0,0 +1,61 @@
+import gzip
+import io
+from dataclasses import dataclass
+
+from .base import CompressionAlgorithm, CompressionResponder, ContentEncoding
+from .types import ASGIApp, Receive, Scope, Send
+
+
+class GzipResponder(CompressionResponder):
+    """Responder that applies gzip compression."""
+
+    content_encoding = ContentEncoding.GZIP
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        minimum_size: int,
+        compresslevel: int = 9,
+    ) -> None:
+        super().__init__(app, minimum_size)
+
+        self.gzip_buffer = io.BytesIO()
+        self.gzip_file = gzip.GzipFile(
+            mode="wb",
+            fileobj=self.gzip_buffer,
+            compresslevel=compresslevel,
+        )
+
+    async def __call__(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+    ) -> None:
+        with self.gzip_buffer, self.gzip_file:
+            await super().__call__(scope, receive, send)
+
+    def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
+        self.gzip_file.write(body)
+        if not more_body:
+            self.gzip_file.close()
+
+        body = self.gzip_buffer.getvalue()
+        self.gzip_buffer.seek(0)
+        self.gzip_buffer.truncate()
+        return body
+
+
+@dataclass
+class GzipAlgorithm(CompressionAlgorithm):
+    """Gzip compression algorithm."""
+
+    type: ContentEncoding = ContentEncoding.GZIP
+    compresslevel: int = 9
+
+    def create_responder(self, app: ASGIApp) -> GzipResponder:
+        return GzipResponder(
+            app=app,
+            minimum_size=self.minimum_size,
+            compresslevel=self.compresslevel,
+        )

asgi_compression/identity.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+from .base import CompressionAlgorithm, CompressionResponder, ContentEncoding
+from .types import ASGIApp
+
+
+class IdentityResponder(CompressionResponder):
+    """Responder that doesn't apply any compression."""
+
+    content_encoding = ContentEncoding.IDENTITY
+
+    def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
+        return body
+
+
+@dataclass
+class IdentityAlgorithm(CompressionAlgorithm):
+    """No compression, identity algorithm."""
+
+    type: ContentEncoding = ContentEncoding.IDENTITY
+
+    def create_responder(self, app: ASGIApp) -> IdentityResponder:
+        return IdentityResponder(app=app, minimum_size=self.minimum_size)

asgi_compression/middleware.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+from .base import (
+    DEFAULT_MINIMUM_SIZE,
+    CompressionAlgorithm,
+    CompressionResponder,
+)
+from .identity import IdentityAlgorithm
+from .types import ASGIApp, Headers, Receive, Scope, Send
+
+
+class CompressionMiddleware:
+    """
+    Unified ASGI middleware for response compression.
+
+    Supports multiple compression algorithms and automatically negotiates
+    the best available algorithm based on the client's Accept-Encoding header.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        algorithms: Optional[List[CompressionAlgorithm]] = None,
+        minimum_size: int = DEFAULT_MINIMUM_SIZE,
+    ) -> None:
+        """
+        Initialize the compression middleware.
+
+        Args:
+            app: The ASGI application.
+            algorithms: List of compression algorithms to use, in order of preference.
+                If not provided, no compression will be applied.
+            minimum_size: The minimum response size to apply compression.
+                This will be used as the default for algorithms that don't specify it.
+        """
+
+        self.app = app
+        self.minimum_size = minimum_size
+
+        self.algorithms = algorithms or []
+        for algorithm in self.algorithms:
+            try:
+                algorithm.check_available()
+            except ImportError as e:
+                raise e from e
+
+        self._default_algorithm = IdentityAlgorithm(minimum_size=minimum_size)
+
+        # Set minimum_size if not explicitly set in the algorithm
+        for algorithm in self.algorithms:
+            if (
+                algorithm.minimum_size == DEFAULT_MINIMUM_SIZE
+                and minimum_size != DEFAULT_MINIMUM_SIZE
+            ):
+                algorithm.minimum_size = minimum_size
+
+    async def __call__(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+    ) -> None:
+        """ASGI application interface."""
+        if scope["type"] != "http":  # pragma: no cover
+            await self.app(scope, receive, send)
+            return
+
+        headers = Headers(scope=scope)
+        accept_encoding = headers.get("Accept-Encoding", "")
+
+        # Find the first supported algorithm that matches the Accept-Encoding header
+        responder: Union[CompressionResponder, None] = None
+        for algorithm in self.algorithms:
+            if str(algorithm.type.value) in accept_encoding:
+                responder = algorithm.create_responder(self.app)
+                break
+
+        # If no matching algorithm, use identity (no compression)
+        if responder is None:
+            responder = self._default_algorithm.create_responder(self.app)
+
+        await responder(scope, receive, send)

asgi_compression/types.py

@@ -0,0 +1,63 @@
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+)
+
+from multidict import CIMultiDict
+
+Scope = MutableMapping[str, Any]
+Message = MutableMapping[str, Any]
+Receive = Callable[[], Awaitable[Message]]
+Send = Callable[[Message], Awaitable[None]]
+ASGIApp = Callable[[Scope, Receive, Send], Awaitable[None]]
+
+
+class Headers(CIMultiDict[str]):
+    def __init__(
+        self,
+        headers: Optional[Mapping[str, str]] = None,
+        raw: Optional[list[tuple[bytes, bytes]]] = None,
+        scope: Optional[Scope] = None,
+    ) -> None:
+        headers_list: List[tuple[str, str]] = []
+        if headers is not None:
+            assert raw is None, 'Cannot set both "headers" and "raw".'
+            assert scope is None, 'Cannot set both "headers" and "scope".'
+            headers_list = list(headers.items())
+        elif raw is not None:
+            assert scope is None, 'Cannot set both "raw" and "scope".'
+            headers_list = [
+                (key.decode("latin-1"), value.decode("latin-1"))
+                for key, value in raw
+            ]
+        elif scope is not None:
+            # scope["headers"] isn't necessarily a list
+            # it might be a tuple or other iterable
+            scope_headers = scope["headers"] = list(scope["headers"])
+            headers_list = [
+                (key.decode("latin-1"), value.decode("latin-1"))
+                for key, value in scope_headers
+            ]
+
+        super().__init__(headers_list)
+
+    def add_vary_header(self, vary: str) -> None:
+        existing = self.get("vary")
+        if existing is not None:
+            # Check if the value is already in the Vary header to avoid duplication
+            values = [x.strip() for x in existing.split(",")]
+            if vary not in values:
+                vary = f"{existing}, {vary}"
+
+        self["vary"] = vary
+
+    def encode(self) -> list[tuple[bytes, bytes]]:
+        return [
+            (key.encode("latin-1"), value.encode("latin-1"))
+            for key, value in self.items()
+        ]

asgi_compression/zstd.py

@@ -0,0 +1,92 @@
+import io
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from .base import CompressionAlgorithm, CompressionResponder, ContentEncoding
+from .types import ASGIApp, Receive, Scope, Send
+
+if TYPE_CHECKING:
+    import zstandard
+
+
+def import_zstandard() -> None:
+    global zstandard
+    try:
+        import zstandard
+    except ImportError as e:
+        raise ImportError(
+            "zstandard is not installed, run `pip install zstandard`]"
+        ) from e
+
+
+class ZstdResponder(CompressionResponder):
+    """Responder that applies Zstandard compression."""
+
+    content_encoding = ContentEncoding.ZSTD
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        minimum_size: int,
+        level: int = 3,
+        threads: int = 0,
+        write_checksum: bool = False,
+        write_content_size: bool = True,
+    ) -> None:
+        super().__init__(app, minimum_size)
+
+        import_zstandard()
+
+        self.zstd_buffer = io.BytesIO()
+        self.compressor = zstandard.ZstdCompressor(
+            level=level,
+            threads=threads,
+            write_checksum=write_checksum,
+            write_content_size=write_content_size,
+        )
+        self.compression_stream = self.compressor.stream_writer(
+            self.zstd_buffer
+        )
+
+    async def __call__(
+        self,
+        scope: Scope,
+        receive: Receive,
+        send: Send,
+    ) -> None:
+        with self.zstd_buffer, self.compression_stream:
+            await super().__call__(scope, receive, send)
+
+    def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
+        self.compression_stream.write(body)
+        if not more_body:
+            self.compression_stream.flush(zstandard.FLUSH_FRAME)
+
+        body = self.zstd_buffer.getvalue()
+        self.zstd_buffer.seek(0)
+        self.zstd_buffer.truncate()
+        return body
+
+
+@dataclass
+class ZstdAlgorithm(CompressionAlgorithm):
+    """Zstandard compression algorithm."""
+
+    type: ContentEncoding = ContentEncoding.ZSTD
+    level: int = 3
+    threads: int = 0
+    write_checksum: bool = False
+    write_content_size: bool = True
+
+    def create_responder(self, app: ASGIApp) -> ZstdResponder:
+        return ZstdResponder(
+            app=app,
+            minimum_size=self.minimum_size,
+            level=self.level,
+            threads=self.threads,
+            write_checksum=self.write_checksum,
+            write_content_size=self.write_content_size,
+        )
+
+    def check_available(self) -> None:
+        import_zstandard()

asgi_compression-0.1.0.dist-info/METADATA

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: asgi-compression
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: multidict>=6.2.0
+Provides-Extra: all
+Requires-Dist: brotli>=1.1.0; extra == "all"
+Requires-Dist: zstandard>=0.23.0; extra == "all"
+Provides-Extra: br
+Requires-Dist: brotli>=1.1.0; extra == "br"
+Provides-Extra: zstd
+Requires-Dist: zstandard>=0.23.0; extra == "zstd"

asgi_compression-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+asgi_compression/__init__.py,sha256=hp71J4CeuSECF52-hVS-mia_m0zcXlHV72qABNww_yY,457
+asgi_compression/base.py,sha256=SmWJFT1Iy5I6Sa4_Vp2uMIr4iaQE6iq686havVoZPmQ,5145
+asgi_compression/brotli.py,sha256=e2YHx1anKJkvHgw0OoDJK_adzZ1ZKLwdrPgDjkUmoQU,2633
+asgi_compression/gzip.py,sha256=dj7kaH5_VbElH2U88oKVcNQPqKTgXl2lE6ghH5Qo0CU,1608
+asgi_compression/identity.py,sha256=tIGUAXxA6YJS0KsZO8se4dlR1_q6_YJNXIBU06-AE98,692
+asgi_compression/middleware.py,sha256=bbH_bdsB3MSgdSN3qMj2VjkGqeGiBGBQRclEsQ6dFHo,2737
+asgi_compression/types.py,sha256=QozTp5Ebehh-I-tv_DkZnjQJ0Z91_Mk5W8-rBI4E0rE,2087
+asgi_compression/zstd.py,sha256=HYX80EzHKy38w5ygAeEcR4Y8a4KVi1UNyns63ziAuvc,2518
+asgi_compression-0.1.0.dist-info/METADATA,sha256=DhbHj5rGQfMI4olwnY3vtrIuJSYq5SyXohPztAf51MI,438
+asgi_compression-0.1.0.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+asgi_compression-0.1.0.dist-info/top_level.txt,sha256=ReOgAA8xTuoMxA4oROBnjJb7yoIjMe7xAB9dz_eXU0Q,17
+asgi_compression-0.1.0.dist-info/RECORD,,

asgi_compression-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (78.1.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

asgi_compression-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+asgi_compression