summaryrefslogtreecommitdiff
path: root/venv/lib/python3.11/site-packages/litestar/middleware/compression/brotli_facade.py
blob: 3d01950a456649ffcb6bf43b0c723d99e5e55742 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
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())