venv

10 files changed, 325 insertions, 0 deletions

diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__init__.py b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__init__.py
new file mode 100644
index 0000000..0885932
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__init__.py
@@ -0,0 +1,4 @@
+from litestar.middleware.compression.facade import CompressionFacade
+from litestar.middleware.compression.middleware import CompressionMiddleware
+
+__all__ = ("CompressionMiddleware", "CompressionFacade")
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/__init__.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000..80ea058
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/__init__.cpython-311.pyc
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/brotli_facade.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/brotli_facade.cpython-311.pyc
new file mode 100644
index 0000000..7378c0f
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/brotli_facade.cpython-311.pyc
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/facade.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/facade.cpython-311.pyc
new file mode 100644
index 0000000..d336c8f
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/facade.cpython-311.pyc
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/gzip_facade.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/gzip_facade.cpython-311.pyc
new file mode 100644
index 0000000..66e1df4
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/gzip_facade.cpython-311.pyc
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/middleware.cpython-311.pyc b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/middleware.cpython-311.pyc
new file mode 100644
index 0000000..a683673
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/__pycache__/middleware.cpython-311.pyc
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/brotli_facade.py b/venv/lib/python3.11/site-packages/litestar/middleware/compression/brotli_facade.py
new file mode 100644
index 0000000..3d01950
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/brotli_facade.py
@@ -0,0 +1,51 @@
+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:
+    from brotli import MODE_FONT, MODE_GENERIC, MODE_TEXT, Compressor
+except ImportError as e:
+    raise MissingDependencyException("brotli") from e
+
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from litestar.config.compression import CompressionConfig
+
+
+class BrotliCompression(CompressionFacade):
+    __slots__ = ("compressor", "buffer", "compression_encoding")
+
+    encoding = CompressionEncoding.BROTLI
+
+    def __init__(
+        self,
+        buffer: BytesIO,
+        compression_encoding: Literal[CompressionEncoding.BROTLI] | str,
+        config: CompressionConfig,
+    ) -> None:
+        self.buffer = buffer
+        self.compression_encoding = compression_encoding
+        modes: dict[Literal["generic", "text", "font"], int] = {
+            "text": int(MODE_TEXT),
+            "font": int(MODE_FONT),
+            "generic": int(MODE_GENERIC),
+        }
+        self.compressor = Compressor(
+            quality=config.brotli_quality,
+            mode=modes[config.brotli_mode],
+            lgwin=config.brotli_lgwin,
+            lgblock=config.brotli_lgblock,
+        )
+
+    def write(self, body: bytes) -> None:
+        self.buffer.write(self.compressor.process(body))
+        self.buffer.write(self.compressor.flush())
+
+    def close(self) -> None:
+        self.buffer.write(self.compressor.finish())
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/facade.py b/venv/lib/python3.11/site-packages/litestar/middleware/compression/facade.py
new file mode 100644
index 0000000..0074b57
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/facade.py
@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar, Protocol
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from litestar.config.compression import CompressionConfig
+    from litestar.enums import CompressionEncoding
+
+
+class CompressionFacade(Protocol):
+    """A unified facade offering a uniform interface for different compression libraries."""
+
+    encoding: ClassVar[str]
+    """The encoding of the compression."""
+
+    def __init__(
+        self, buffer: BytesIO, compression_encoding: CompressionEncoding | str, config: CompressionConfig
+    ) -> None:
+        """Initialize ``CompressionFacade``.
+
+        Args:
+            buffer: A bytes IO buffer to write the compressed data into.
+            compression_encoding: The compression encoding used.
+            config: The app compression config.
+        """
+        ...
+
+    def write(self, body: bytes) -> None:
+        """Write compressed bytes.
+
+        Args:
+            body: Message body to process
+
+        Returns:
+            None
+        """
+        ...
+
+    def close(self) -> None:
+        """Close the compression stream.
+
+        Returns:
+            None
+        """
+        ...
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/gzip_facade.py b/venv/lib/python3.11/site-packages/litestar/middleware/compression/gzip_facade.py
new file mode 100644
index 0000000..b10ef73
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/gzip_facade.py
@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from gzip import GzipFile
+from typing import TYPE_CHECKING, Literal
+
+from litestar.enums import CompressionEncoding
+from litestar.middleware.compression.facade import CompressionFacade
+
+if TYPE_CHECKING:
+    from io import BytesIO
+
+    from litestar.config.compression import CompressionConfig
+
+
+class GzipCompression(CompressionFacade):
+    __slots__ = ("compressor", "buffer", "compression_encoding")
+
+    encoding = CompressionEncoding.GZIP
+
+    def __init__(
+        self, buffer: BytesIO, compression_encoding: Literal[CompressionEncoding.GZIP] | str, config: CompressionConfig
+    ) -> None:
+        self.buffer = buffer
+        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)
+        self.compressor.flush()
+
+    def close(self) -> None:
+        self.compressor.close()
diff --git a/venv/lib/python3.11/site-packages/litestar/middleware/compression/middleware.py b/venv/lib/python3.11/site-packages/litestar/middleware/compression/middleware.py
new file mode 100644
index 0000000..7ea7853
--- /dev/null
+++ b/venv/lib/python3.11/site-packages/litestar/middleware/compression/middleware.py
@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Literal
+
+from litestar.datastructures import Headers, MutableScopeHeaders
+from litestar.enums import CompressionEncoding, ScopeType
+from litestar.middleware.base import AbstractMiddleware
+from litestar.middleware.compression.gzip_facade import GzipCompression
+from litestar.utils.empty import value_or_default
+from litestar.utils.scope.state import ScopeState
+
+if TYPE_CHECKING:
+    from litestar.config.compression import CompressionConfig
+    from litestar.middleware.compression.facade import CompressionFacade
+    from litestar.types import (
+        ASGIApp,
+        HTTPResponseStartEvent,
+        Message,
+        Receive,
+        Scope,
+        Send,
+    )
+
+    try:
+        from brotli import Compressor
+    except ImportError:
+        Compressor = Any
+
+
+class CompressionMiddleware(AbstractMiddleware):
+    """Compression Middleware Wrapper.
+
+    This is a wrapper allowing for generic compression configuration / handler middleware
+    """
+
+    def __init__(self, app: ASGIApp, config: CompressionConfig) -> None:
+        """Initialize ``CompressionMiddleware``
+
+        Args:
+            app: The ``next`` ASGI app to call.
+            config: An instance of CompressionConfig.
+        """
+        super().__init__(
+            app=app, exclude=config.exclude, exclude_opt_key=config.exclude_opt_key, scopes={ScopeType.HTTP}
+        )
+        self.config = config
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """ASGI callable.
+
+        Args:
+            scope: The ASGI connection scope.
+            receive: The ASGI receive function.
+            send: The ASGI send function.
+
+        Returns:
+            None
+        """
+        accept_encoding = Headers.from_scope(scope).get("accept-encoding", "")
+        config = self.config
+
+        if config.compression_facade.encoding in accept_encoding:
+            await self.app(
+                scope,
+                receive,
+                self.create_compression_send_wrapper(
+                    send=send, compression_encoding=config.compression_facade.encoding, scope=scope
+                ),
+            )
+            return
+
+        if config.gzip_fallback and CompressionEncoding.GZIP in accept_encoding:
+            await self.app(
+                scope,
+                receive,
+                self.create_compression_send_wrapper(
+                    send=send, compression_encoding=CompressionEncoding.GZIP, scope=scope
+                ),
+            )
+            return
+
+        await self.app(scope, receive, send)
+
+    def create_compression_send_wrapper(
+        self,
+        send: Send,
+        compression_encoding: Literal[CompressionEncoding.BROTLI, CompressionEncoding.GZIP] | str,
+        scope: Scope,
+    ) -> Send:
+        """Wrap ``send`` to handle brotli compression.
+
+        Args:
+            send: The ASGI send function.
+            compression_encoding: The compression encoding used.
+            scope: The ASGI connection scope
+
+        Returns:
+            An ASGI send function.
+        """
+        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
+            )
+
+        initial_message: HTTPResponseStartEvent | None = None
+        started = False
+
+        connection_state = ScopeState.from_scope(scope)
+
+        async def send_wrapper(message: Message) -> None:
+            """Handle and compresses the HTTP Message with brotli.
+
+            Args:
+                message (Message): An ASGI Message.
+            """
+            nonlocal started
+            nonlocal initial_message
+
+            if message["type"] == "http.response.start":
+                initial_message = message
+                return
+
+            if initial_message is not None and value_or_default(connection_state.is_cached, False):
+                await send(initial_message)
+                await send(message)
+                return
+
+            if initial_message and message["type"] == "http.response.body":
+                body = message["body"]
+                more_body = message.get("more_body")
+
+                if not started:
+                    started = True
+                    if more_body:
+                        headers = MutableScopeHeaders(initial_message)
+                        headers["Content-Encoding"] = compression_encoding
+                        headers.extend_header_value("vary", "Accept-Encoding")
+                        del headers["Content-Length"]
+                        connection_state.response_compressed = True
+
+                        facade.write(body)
+
+                        message["body"] = bytes_buffer.getvalue()
+                        bytes_buffer.seek(0)
+                        bytes_buffer.truncate()
+                        await send(initial_message)
+                        await send(message)
+
+                    elif len(body) >= self.config.minimum_size:
+                        facade.write(body)
+                        facade.close()
+                        body = bytes_buffer.getvalue()
+
+                        headers = MutableScopeHeaders(initial_message)
+                        headers["Content-Encoding"] = compression_encoding
+                        headers["Content-Length"] = str(len(body))
+                        headers.extend_header_value("vary", "Accept-Encoding")
+                        message["body"] = body
+                        connection_state.response_compressed = True
+
+                        await send(initial_message)
+                        await send(message)
+
+                    else:
+                        await send(initial_message)
+                        await send(message)
+
+                else:
+                    facade.write(body)
+                    if not more_body:
+                        facade.close()
+
+                    message["body"] = bytes_buffer.getvalue()
+
+                    bytes_buffer.seek(0)
+                    bytes_buffer.truncate()
+
+                    if not more_body:
+                        bytes_buffer.close()
+
+                    await send(message)
+
+        return send_wrapper