Spaces:
Running
Running
File size: 6,047 Bytes
e9efb20 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | from __future__ import annotations
import gzip
import io
from typing import NoReturn
from starlette.datastructures import Headers, MutableHeaders
from starlette.types import ASGIApp, Message, Receive, Scope, Send
try:
import brotli
HAS_BROTLI = True
except ImportError:
brotli = None
HAS_BROTLI = False
COMPRESSIBLE_CONTENT_TYPES = (
"application/json",
"application/javascript",
"application/manifest+json",
"application/xml",
"application/xhtml+xml",
"image/svg+xml",
"text/",
)
class CompressionMiddleware:
"""Negotiates Brotli, then gzip, then identity based on Accept-Encoding.
Only text-like content types are compressed, so already-compressed media
(images, audio, video) and zero-copy file sends pass through untouched.
"""
def __init__(
self,
app: ASGIApp,
minimum_size: int = 512,
gzip_level: int = 6,
brotli_quality: int = 5,
) -> None:
self.app = app
self.minimum_size = minimum_size
self.gzip_level = gzip_level
self.brotli_quality = brotli_quality
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
accept_encoding = Headers(scope=scope).get("Accept-Encoding", "")
responder: _CompressionResponder
if HAS_BROTLI and "br" in accept_encoding:
responder = _BrotliResponder(
self.app, self.minimum_size, self.brotli_quality
)
elif "gzip" in accept_encoding:
responder = _GZipResponder(self.app, self.minimum_size, self.gzip_level)
else:
responder = _CompressionResponder(self.app, self.minimum_size)
await responder(scope, receive, send)
class _CompressionResponder:
content_encoding: str = ""
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.passthrough = True
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":
self.initial_message = message
headers = Headers(raw=message["headers"])
content_type = headers.get("content-type", "")
already_encoded = "content-encoding" in headers
self.passthrough = already_encoded or not content_type.startswith(
COMPRESSIBLE_CONTENT_TYPES
)
elif message_type == "http.response.body" and self.passthrough:
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:
await self.send(self.initial_message)
await self.send(message)
else:
compressed = self.apply_compression(body, more_body=more_body)
headers = MutableHeaders(raw=self.initial_message["headers"])
headers.add_vary_header("Accept-Encoding")
if compressed != body:
headers["Content-Encoding"] = self.content_encoding
if more_body:
del headers["Content-Length"]
else:
headers["Content-Length"] = str(len(compressed))
message["body"] = compressed
await self.send(self.initial_message)
await self.send(message)
elif message_type == "http.response.body":
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)
elif message_type == "http.response.pathsend":
await self.send(self.initial_message)
await self.send(message)
def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
return body
class _GZipResponder(_CompressionResponder):
content_encoding = "gzip"
def __init__(self, app: ASGIApp, minimum_size: int, compresslevel: int) -> None:
super().__init__(app, minimum_size)
self.buffer = io.BytesIO()
self.file = gzip.GzipFile(
mode="wb", fileobj=self.buffer, compresslevel=compresslevel
)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
with self.buffer, self.file:
await super().__call__(scope, receive, send)
def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
self.file.write(body)
if not more_body:
self.file.close()
data = self.buffer.getvalue()
self.buffer.seek(0)
self.buffer.truncate()
return data
class _BrotliResponder(_CompressionResponder):
content_encoding = "br"
def __init__(self, app: ASGIApp, minimum_size: int, quality: int) -> None:
super().__init__(app, minimum_size)
self.compressor = brotli.Compressor(quality=quality)
def apply_compression(self, body: bytes, *, more_body: bool) -> bytes:
data = self.compressor.process(body)
if not more_body:
data += self.compressor.finish()
return data
async def _unattached_send(message: Message) -> NoReturn:
raise RuntimeError("send awaitable not set")
|