Repository: Neoteroi/BlackSheep
Pull request #329 (2023-04-17 17:17:05)

Add features to describe app environment


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8ab865e..ddcc170 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [1.2.13] - 2023-04-17 :crown:
+
+- Adds a `GzipMiddleware` that can be used to enable `gzip` compression, using
+  the built-in module. [tyzhnenko](https://github.com/tyzhnenko)'s contribution. :sparkles:
+- Adds a strategy to control features depending on application environment:
+  `is_development`, `is_production` depending on `APP_ENV` environment
+  variable. For more information, see [_Defining application environment_](https://www.neoteroi.dev/blacksheep/settings/#defining-application-environment).
+- Makes the client `ConnectionPools` a context manager, its `__exit__` method
+  closes all its `TCP-IP` connections.
+
 ## [1.2.12] - 2023-03-22 :flamingo:
 
 - Adds `@app.lifespan` to support registering objects that must be initialized
diff --git a/blacksheep/client/pool.py b/blacksheep/client/pool.py
index 320f621..d529e6f 100644
--- a/blacksheep/client/pool.py
+++ b/blacksheep/client/pool.py
@@ -128,6 +128,12 @@ class ConnectionPools:
             self._pools[key] = new_pool
             return new_pool
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.dispose()
+
     def dispose(self):
         for pool in self._pools.values():
             pool.dispose()
diff --git a/blacksheep/server/env.py b/blacksheep/server/env.py
index 4adc2a3..85f5026 100644
--- a/blacksheep/server/env.py
+++ b/blacksheep/server/env.py
@@ -1,16 +1,39 @@
 import os
+from dataclasses import dataclass
 
+from blacksheep.utils import truthy
 
-def truthy(value: str, default: bool = False) -> bool:
-    if not value:
-        return default
-    return value.upper() in {"1", "TRUE"}
 
+def get_env() -> str:
+    return os.environ.get("APP_ENV", "production")
 
+
+def is_development() -> bool:
+    """
+    Returns a value indicating whether the application is running for local development.
+    This method checks if an `APP_ENV` environment variable is set and its lowercase
+    value is either "local", "dev", or "development".
+    """
+    return get_env().lower() in {"local", "dev", "development"}
+
+
+def is_production() -> bool:
+    """
+    Returns a value indicating whether the application is running for the production
+    environment (default is true).
+    This method checks if an `APP_ENV` environment variable is set and its lowercase
+    value is either "prod" or "production".
+    """
+    return get_env().lower() in {"prod", "production"}
+
+
+@dataclass(init=False)
 class EnvironmentSettings:
+    env: str
     show_error_details: bool
     mount_auto_events: bool
 
     def __init__(self) -> None:
+        self.env = get_env()
         self.show_error_details = truthy(os.environ.get("APP_SHOW_ERROR_DETAILS", ""))
         self.mount_auto_events = truthy(os.environ.get("APP_MOUNT_AUTO_EVENTS", ""))
diff --git a/blacksheep/server/gzip.py b/blacksheep/server/gzip.py
new file mode 100644
index 0000000..10fab7b
--- /dev/null
+++ b/blacksheep/server/gzip.py
@@ -0,0 +1,152 @@
+import asyncio
+import gzip
+from concurrent.futures import Executor
+from typing import Awaitable, Callable, Iterable, List, Optional
+
+from blacksheep import Content, Request, Response
+from blacksheep.server.application import Application
+from blacksheep.server.normalization import ensure_response
+
+
+class GzipMiddleware:
+    """
+    The gzip compression middleware for all requests with a body larger than
+    the specified minimum size and with the "gzip" encoding in the "Accept-Encoding"
+    header. The middleware runs compression asynchronously in a separate executor.
+
+    Parameters
+    ----------
+    min_size: int
+        The minimum size of the response body to compress.
+    comp_level: int
+        The compression level to use.
+    handled_types: Optional[Iterable[bytes]]
+        The list of content types to compress.
+    executor: Executor
+        The executor instance to use for compression. If not specified, a
+        default executor is used. If you specify an executor, you are responsible
+        for shutting it down.
+    """
+
+    handled_types: List[bytes] = [
+        b"json",
+        b"xml",
+        b"yaml",
+        b"html",
+        b"text/plain",
+        b"application/javascript",
+        b"text/css",
+        b"text/csv",
+    ]
+
+    def __init__(
+        self,
+        min_size: int = 500,
+        comp_level: int = 5,
+        handled_types: Optional[Iterable[bytes]] = None,
+        executor: Optional[Executor] = None,
+    ):
+        self.min_size = min_size
+        self.comp_level = comp_level
+        self._executor = executor
+
+        if handled_types is not None:
+            self.handled_types = self._normalize_types(handled_types)
+
+    def _normalize_types(self, types: Iterable[bytes]) -> List[bytes]:
+        """
+        Normalizes the types to bytes.
+        """
+        normalized_types = []
+        for _type in types:
+            if isinstance(_type, str):
+                normalized_types.append(_type.encode("ascii"))
+            else:
+                normalized_types.append(_type)
+        return normalized_types
+
+    def should_handle(self, request: Request, response: Response) -> bool:
+        """
+        Returns True if the response should be compressed.
+        """
+
+        def _is_handled_type(content_type) -> bool:
+            content_type = content_type.lower()
+            return any(_type in content_type for _type in self.handled_types)
+
+        def is_handled_encoding() -> bool:
+            return request.headers is not None and b"gzip" in (
+                request.headers.get_single(b"accept-encoding") or ""
+            )
+
+        def is_handled_response_content() -> bool:
+            if response is None or response.content is None:
+                return False
+
+            body_pass: bool = (
+                response.content.body is not None
+                and len(response.content.body) > self.min_size
+            )
+
+            content_type_pass: bool = (
+                response.content.type is not None
+                and _is_handled_type(response.content.type)
+            )
+
+            return all(
+                (
+                    body_pass,
+                    content_type_pass,
+                )
+            )
+
+        return all(
+            (
+                is_handled_encoding(),
+                is_handled_response_content(),
+            )
+        )
+
+    async def __call__(
+        self, request: Request, handler: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        response = ensure_response(await handler(request))
+
+        if response is None or response.content is None:
+            return response
+
+        if not self.should_handle(request, response):
+            return response
+
+        loop = asyncio.get_running_loop()
+        compressed_body = await loop.run_in_executor(
+            self._executor,
+            gzip.compress,
+            response.content.body,
+            self.comp_level,
+        )
+
+        response.with_content(
+            Content(
+                content_type=response.content.type,
+                data=compressed_body,
+            )
+        )
+        response.add_header(b"content-encoding", b"gzip")
+        return response
+
+
+def use_gzip_compression(
+    app: Application,
+    handler: Optional[GzipMiddleware] = None,
+) -> GzipMiddleware:
+    """
+    Configures the application to use gzip compression for all responses with gzip
+    in accept-encoding header.
+    """
+    if handler is None:
+        handler = GzipMiddleware()
+
+    app.middlewares.append(handler)
+
+    return handler
diff --git a/blacksheep/server/security/hsts.py b/blacksheep/server/security/hsts.py
index 46245a0..d3a9a0e 100644
--- a/blacksheep/server/security/hsts.py
+++ b/blacksheep/server/security/hsts.py
@@ -1,7 +1,7 @@
 from blacksheep.messages import Request, Response
 
 
-def _get_hsts_value(max_age: int, include_subdomains: bool) -> bytes:
+def write_hsts_header_value(max_age: int, include_subdomains: bool) -> bytes:
     value = f"max-age={max_age};"
 
     if include_subdomains:
@@ -21,7 +21,7 @@ class HSTSMiddleware:
         max_age: int = 31536000,
         include_subdomains: bool = True,
     ) -> None:
-        self._value = _get_hsts_value(max_age, include_subdomains)
+        self._value = write_hsts_header_value(max_age, include_subdomains)
 
     async def __call__(self, request: Request, handler):
         response: Response = await handler(request)
diff --git a/blacksheep/utils/__init__.py b/blacksheep/utils/__init__.py
index 0f79bd3..59ce661 100644
--- a/blacksheep/utils/__init__.py
+++ b/blacksheep/utils/__init__.py
@@ -37,3 +37,9 @@ def get_class_hierarchy(cls: Type[T]) -> Tuple[Type[T], ...]:
 
 def get_class_instance_hierarchy(instance: T) -> Tuple[Type[T], ...]:
     return get_class_hierarchy(type(instance))
+
+
+def truthy(value: str, default: bool = False) -> bool:
+    if not value:
+        return default
+    return value.upper() in {"1", "TRUE"}
diff --git a/setup.py b/setup.py
index 3df0f82..6acbdef 100644
--- a/setup.py
+++ b/setup.py
@@ -11,7 +11,7 @@ COMPILE_ARGS = ["-O2"]
 
 setup(
     name="blacksheep",
-    version="1.2.12",
+    version="1.2.13",
     description="Fast web framework for Python asyncio",
     long_description=readme(),
     long_description_content_type="text/markdown",


Test patch:
diff --git a/itests/app_1.py b/itests/app_1.py
index 7fe41b5..37f4478 100644
--- a/itests/app_1.py
+++ b/itests/app_1.py
@@ -17,10 +17,12 @@ from blacksheep import (
     json,
     text,
 )
+from blacksheep.server.gzip import use_gzip_compression
 from itests.utils import CrashTest, ensure_folder
 
 app = Application(show_error_details=True)
 
+use_gzip_compression(app)
 
 static_folder_path = pathlib.Path(__file__).parent.absolute() / "static"
 
diff --git a/itests/app_2.py b/itests/app_2.py
index 45da5ae..ca8a131 100644
--- a/itests/app_2.py
+++ b/itests/app_2.py
@@ -31,6 +31,7 @@ from blacksheep.server.bindings import (
     FromServices,
 )
 from blacksheep.server.controllers import APIController
+from blacksheep.server.gzip import use_gzip_compression
 from blacksheep.server.openapi.common import (
     ContentInfo,
     EndpointDocs,
@@ -49,6 +50,7 @@ from itests.utils import CrashTest
 
 app_2 = Application()
 
+use_gzip_compression(app_2)
 
 controllers_router = RoutesRegistry()
 app_2.controllers_router = controllers_router
diff --git a/itests/app_3.py b/itests/app_3.py
index 7c23218..849b7e8 100644
--- a/itests/app_3.py
+++ b/itests/app_3.py
@@ -1,11 +1,14 @@
 import uvicorn
 
 from blacksheep.server import Application
+from blacksheep.server.gzip import use_gzip_compression
 from blacksheep.server.responses import json
 
 application = Application(show_error_details=True)
 app_3 = Application(show_error_details=True)
 
+use_gzip_compression(app_3)
+
 
 @application.router.get("/foo")
 def handle_foo():
diff --git a/itests/app_4.py b/itests/app_4.py
index 13aa429..b0373b6 100644
--- a/itests/app_4.py
+++ b/itests/app_4.py
@@ -11,6 +11,7 @@ from blacksheep.plugins import json as json_plugin
 from blacksheep.plugins.json import default_json_dumps
 from blacksheep.server import Application
 from blacksheep.server.bindings import FromJSON
+from blacksheep.server.gzip import use_gzip_compression
 from blacksheep.server.responses import json
 from blacksheep.server.websocket import WebSocket
 
@@ -70,6 +71,8 @@ def configure_json_settings():
 
 app_4 = Application(show_error_details=True)
 
+use_gzip_compression(app_4)
+
 
 @dataclasses.dataclass
 class MyData:
diff --git a/tests/test_gzip.py b/tests/test_gzip.py
new file mode 100644
index 0000000..5781387
--- /dev/null
+++ b/tests/test_gzip.py
@@ -0,0 +1,129 @@
+import gzip
+
+import pytest
+
+from blacksheep.server.gzip import GzipMiddleware
+from blacksheep.testing.helpers import get_example_scope
+from blacksheep.testing.messages import MockReceive, MockSend
+
+
+@pytest.mark.parametrize(
+    "comp_level, comp_size",
+    (zip(range(0, 10), (468, 283, 283, 283, 282, 282, 282, 282, 282, 282))),
+)
+@pytest.mark.asyncio
+async def test_gzip_output(app, comp_level, comp_size):
+    return_value = (
+        "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do "
+        "eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim "
+        "veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo "
+        "consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum "
+        "dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, "
+        "sunt in culpa qui officia deserunt mollit anim id est laborum."
+    )
+
+    @app.router.get("/")
+    async def home():
+        return return_value
+
+    app.middlewares.append(GzipMiddleware(min_size=0, comp_level=comp_level))
+
+    await app.start()
+
+    await app(
+        get_example_scope(
+            "GET",
+            "/",
+        ),
+        MockReceive([]),
+        MockSend(),
+    )
+
+    response = app.response
+    assert response.status == 200
+    assert response.content.length == comp_size
+    assert gzip.decompress(response.content.body) == return_value.encode("ascii")
+    assert response.headers.get_single(b"content-encoding") == b"gzip"
+
+
+@pytest.mark.asyncio
+async def test_skip_gzip_small_output(app):
+    @app.router.get("/")
+    async def home():
+        return "Hello, World"
+
+    app.middlewares.append(GzipMiddleware(min_size=16))
+
+    await app.start()
+
+    await app(
+        get_example_scope(
+            "GET",
+            "/",
+        ),
+        MockReceive([]),
+        MockSend(),
+    )
+
+    response = app.response
+    assert response.status == 200
+    assert response.content.body == b"Hello, World"
+    assert response.content.length == 12
+    with pytest.raises(ValueError):
+        assert response.headers.get_single(b"content-encoding")
+
+
+@pytest.mark.asyncio
+async def test_skip_gzip_output_without_header(app):
+    @app.router.get("/")
+    async def home():
+        return "Hello, World"
+
+    app.middlewares.append(GzipMiddleware(min_size=0))
+
+    await app.start()
+
+    await app(
+        get_example_scope(
+            "GET",
+            "/",
+            accept_encoding=b"deflate",
+        ),
+        MockReceive([]),
+        MockSend(),
+    )
+
+    response = app.response
+    assert response.status == 200
+    assert response.content.body == b"Hello, World"
+    assert response.content.length == 12
+    with pytest.raises(ValueError):
+        assert response.headers.get_single(b"content-encoding")
+
+
+@pytest.mark.asyncio
+async def test_skip_gzip_output_for_unhandled_type(app):
+    @app.router.get("/")
+    async def home():
+        return "Hello, World"
+
+    app.middlewares.append(GzipMiddleware(min_size=0, handled_types=[b"text/html"]))
+
+    await app.start()
+
+    await app(
+        get_example_scope(
+            "GET",
+            "/",
+            accept_encoding=b"deflate",
+        ),
+        MockReceive([]),
+        MockSend(),
+    )
+
+    response = app.response
+    assert response.status == 200
+    assert response.content.length == 12
+    assert response.content.body == b"Hello, World"
+    with pytest.raises(ValueError):
+        assert response.headers.get_single(b"content-encoding")