From 41258066a7435441b6995a9418acc6dce16469b4 Mon Sep 17 00:00:00 2001 From: dbf <13336073+fullonic@users.noreply.github.com> Date: Wed, 28 Sep 2022 13:18:04 +0200 Subject: [PATCH] Feature: Allow excluding handlers from being compressed (#21) * Update README * Add `excluded_handlers` params to middleware * Use `typing.NoReturn` instead of `None` * Update docstring * Refactor `BrotliMiddleware.__call__` --- .gitignore | 1 + README.md | 2 +- brotli_asgi/__init__.py | 65 +++++++++++++++++++++++++---------------- tests.py | 27 ++++++++++++++++- 4 files changed, 68 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index 9d21c80..7bfae7b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.venv venv .venv __pycache__ diff --git a/README.md b/README.md index 18c300e..56974cb 100644 --- a/README.md +++ b/README.md @@ -100,4 +100,4 @@ sys.getsizeof(gzip.compress(page, compresslevel=6)) ## Compatibility -According to [caniuse.com](https://caniuse.com/#feat=brotli), Brotli is supported by all major browsers with a global use of over _95%_. +According to [caniuse.com](https://caniuse.com/#feat=brotli), Brotli is supported by all major browsers with a global use of over _96.3%_. diff --git a/brotli_asgi/__init__.py b/brotli_asgi/__init__.py index 91b354e..d0f5e64 100644 --- a/brotli_asgi/__init__.py +++ b/brotli_asgi/__init__.py @@ -4,6 +4,8 @@ """ import io +import re +from typing import List, Union, NoReturn from brotli import MODE_FONT, MODE_GENERIC, MODE_TEXT, Compressor # type: ignore from starlette.datastructures import Headers, MutableHeaders @@ -31,7 +33,8 @@ def __init__( lgblock: int = 0, minimum_size: int = 400, gzip_fallback: bool = True, - ) -> None: + excluded_handlers: Union[List, None] = None, + ) -> NoReturn: """ Arguments. @@ -48,6 +51,7 @@ def __init__( quality. minimum_size: Only compress responses that are bigger than this value in bytes. gzip_fallback: If True, uses gzip encoding if br is not in the Accept-Encoding header. + excluded_handlers: List of handlers to be excluded from being compressed. """ self.app = app self.quality = quality @@ -56,27 +60,37 @@ def __init__( self.lgwin = lgwin self.lgblock = lgblock self.gzip_fallback = gzip_fallback - - async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: - if scope["type"] == "http": - headers = Headers(scope=scope) - if "br" in headers.get("Accept-Encoding", ""): - responder = BrotliResponder( - self.app, - self.quality, - self.mode, - self.lgwin, - self.lgblock, - self.minimum_size, - ) - await responder(scope, receive, send) - return - if self.gzip_fallback and "gzip" in headers.get("Accept-Encoding", ""): - responder = GZipResponder(self.app, self.minimum_size) - await responder(scope, receive, send) - return + if excluded_handlers: + self.excluded_handlers = [re.compile(path) for path in excluded_handlers] + else: + self.excluded_handlers = [] + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> NoReturn: + if self._is_handler_excluded(scope) or scope["type"] != "http": + return await self.app(scope, receive, send) + headers = Headers(scope=scope) + if "br" in headers.get("Accept-Encoding", ""): + responder = BrotliResponder( + self.app, + self.quality, + self.mode, + self.lgwin, + self.lgblock, + self.minimum_size, + ) + await responder(scope, receive, send) + return + if self.gzip_fallback and "gzip" in headers.get("Accept-Encoding", ""): + responder = GZipResponder(self.app, self.minimum_size) + await responder(scope, receive, send) + return await self.app(scope, receive, send) + def _is_handler_excluded(self, scope: Scope) -> bool: + handler = scope.get("path", "") + + return any(pattern.search(handler) for pattern in self.excluded_handlers) + class BrotliResponder: """Brotli Interface.""" @@ -89,7 +103,7 @@ def __init__( lgwin: int, lgblock: int, minimum_size: int, - ) -> None: # noqa + ) -> NoReturn: # noqa self.app = app self.quality = quality self.mode = mode @@ -106,11 +120,11 @@ def __init__( async def __call__( self, scope: Scope, receive: Receive, send: Send - ) -> None: # noqa + ) -> NoReturn: # noqa self.send = send await self.app(scope, receive, self.send_with_brotli) - async def send_with_brotli(self, message: Message) -> None: + async def send_with_brotli(self, message: Message) -> NoReturn: """Apply compression using brotli.""" message_type = message["type"] if message_type == "http.response.start": @@ -173,10 +187,11 @@ def _process(self, body): identical except that the official Google API has Compressor.process while the brotlipy API has Compress.compress """ - if hasattr(self.br_file, 'process'): + if hasattr(self.br_file, "process"): return self.br_file.process(body) return self.br_file.compress(body) -async def unattached_send(message: Message) -> None: + +async def unattached_send(message: Message) -> NoReturn: raise RuntimeError("send awaitable not set") # pragma: no cover diff --git a/tests.py b/tests.py index 7ff744d..dbfab57 100644 --- a/tests.py +++ b/tests.py @@ -90,7 +90,11 @@ def test_brotli_api_options(): app = Starlette() app.add_middleware( - BrotliMiddleware, quality=11, mode="text", lgwin=20, lgblock=16, + BrotliMiddleware, + quality=11, + mode="text", + lgwin=20, + lgblock=16, ) @app.route("/") @@ -134,3 +138,24 @@ def homepage(request): assert response.text == "x" * 4000 assert "Content-Encoding" not in response.headers assert int(response.headers["Content-Length"]) == 4000 + + +def test_excluded_handlers(): + app = Starlette() + + app.add_middleware( + BrotliMiddleware, + excluded_handlers=["/excluded"], + ) + + @app.route("/excluded") + def homepage(request): + return PlainTextResponse("x" * 4000, status_code=200) + + client = TestClient(app) + response = client.get("/excluded", headers={"accept-encoding": "br"}) + + assert response.status_code == 200 + assert response.text == "x" * 4000 + assert "Content-Encoding" not in response.headers + assert int(response.headers["Content-Length"]) == 4000