Merge pull request #10 from taskiq-python/bugfix/swagger

s3rius · web-flow · commit 8ceab26f55fc · 2023-07-25T22:10:19.000+04:00
diff --git a/README.md b/README.md
@@ -69,8 +69,10 @@ main_router.add_routes(memes_router, prefix="/memes")
 
 If you use dependencies in you handlers, we can easily generate swagger for you.
 We have some limitations:
-1. We don't support string type annotation for detecting required parameters in openapi. Like `a: "Optional[int]"`.
-2. We don't have support for 3.10 style Option annotations. E.G. `int | None`
+1. We don't support resolving type aliases if hint is a string.
+    If you define variable like this: `myvar = int | None` and then in handler
+    you'd create annotation like this: `param: "str | myvar"` it will fail.
+    You need to unquote type hint in order to get it work.
 
 We will try to fix these limitations later.
 
@@ -84,6 +86,46 @@ app = web.Application()
 app.on_startup.extend([init, setup_swagger()])
 ```
 
+### Responses
+
+You can define schema for responses using dataclasses or
+pydantic models. This would not affect handlers in any way,
+it's only for documentation purposes, if you want to actually
+validate values your handler returns, please write your own wrapper.
+
+```python
+from dataclasses import dataclass
+
+from aiohttp import web
+from pydantic import BaseModel
+
+from aiohttp_deps import Router, openapi_response
+
+router = Router()
+
+
+@dataclass
+class Success:
+    data: str
+
+
+class Unauthorized(BaseModel):
+    why: str
+
+
+@router.get("/")
+@openapi_response(200, Success, content_type="application/xml")
+@openapi_response(200, Success)
+@openapi_response(401, Unauthorized, description="When token is not correct")
+async def handler() -> web.Response:
+    ...
+```
+
+This example illustrates how much you can do with this decorator. You
+can have multiple content-types for a single status, or you can have different
+possble statuses. This function is pretty simple and if you want to make
+your own decorator for your responses, it won't be hard.
+
 
 ## Default dependencies
 
diff --git a/aiohttp_deps/__init__.py b/aiohttp_deps/__init__.py
@@ -3,7 +3,7 @@
 
 from aiohttp_deps.initializer import init
 from aiohttp_deps.router import Router
-from aiohttp_deps.swagger import extra_openapi, setup_swagger
+from aiohttp_deps.swagger import extra_openapi, openapi_response, setup_swagger
 from aiohttp_deps.utils import Form, Header, Json, Path, Query
 from aiohttp_deps.view import View
 
@@ -19,4 +19,5 @@
     "Query",
     "Form",
     "Path",
+    "openapi_response",
 ]
diff --git a/aiohttp_deps/swagger.py b/aiohttp_deps/swagger.py
@@ -1,17 +1,18 @@
 import inspect
 from collections import defaultdict
 from logging import getLogger
-from typing import Any, Awaitable, Callable, Dict, Optional, Union
+from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, get_type_hints
 
 import pydantic
 from aiohttp import web
-from pydantic.utils import deep_update
+from deepmerge import always_merger
 from taskiq_dependencies import DependencyGraph
 
 from aiohttp_deps.initializer import InjectableFuncHandler, InjectableViewHandler
 from aiohttp_deps.utils import Form, Header, Json, Path, Query
 
-REF_TEMPLATE = "#/components/schemas/{model}"
+_T = TypeVar("_T")  # noqa: WPS111
+
 SCHEMA_KEY = "openapi_schema"
 SWAGGER_HTML_TEMPALTE = """
 <html lang="en">
@@ -67,19 +68,14 @@ def _is_optional(annotation: Optional[inspect.Parameter]) -> bool:
     if annotation is None or annotation.annotation == annotation.empty:
         return True
 
-    origin = getattr(annotation.annotation, "__origin__", None)
-    if origin is None:
-        return False
+    def dummy(_var: annotation.annotation) -> None:  # type: ignore
+        """Dummy function to use for type resolution."""
 
-    if origin == Union:
-        args = getattr(annotation.annotation, "__args__", ())
-        for arg in args:
-            if arg is type(None):  # noqa: E721, WPS516
-                return True
-    return False
+    var = get_type_hints(dummy).get("_var")
+    return var == Optional[var]
 
 
-def _add_route_def(  # noqa: C901
+def _add_route_def(  # noqa: C901, WPS210
     openapi_schema: Dict[str, Any],
     route: web.ResourceRoute,
     method: str,
@@ -94,6 +90,19 @@ def _add_route_def(  # noqa: C901
     if route.resource is None:  # pragma: no cover
         return
 
+    params: Dict[tuple[str, str], Any] = {}
+
+    def _insert_in_params(data: Dict[str, Any]) -> None:
+        element = params.get((data["name"], data["in"]))
+        if element is None:
+            params[(data["name"], data["in"])] = data
+            return
+        element["required"] = element.get("required") or data.get("required")
+        element["allowEmptyValue"] = bool(element.get("allowEmptyValue")) and bool(
+            data.get("allowEmptyValue"),
+        )
+        params[(data["name"], data["in"])] = element
+
     for dependency in graph.ordered_deps:
         if isinstance(dependency.dependency, (Json, Form)):
             content_type = "application/json"
@@ -105,9 +114,7 @@ def _add_route_def(  # noqa: C901
             ):
                 input_schema = pydantic.TypeAdapter(
                     dependency.signature.annotation,
-                ).json_schema(
-                    ref_template=REF_TEMPLATE,
-                )
+                ).json_schema()
                 openapi_schema["components"]["schemas"].update(
                     input_schema.pop("definitions", {}),
                 )
@@ -119,7 +126,7 @@ def _add_route_def(  # noqa: C901
                     "content": {content_type: {}},
                 }
         elif isinstance(dependency.dependency, Query):
-            route_info["parameters"].append(
+            _insert_in_params(
                 {
                     "name": dependency.dependency.alias or dependency.param_name,
                     "in": "query",
@@ -128,16 +135,17 @@ def _add_route_def(  # noqa: C901
                 },
             )
         elif isinstance(dependency.dependency, Header):
-            route_info["parameters"].append(
+            name = dependency.dependency.alias or dependency.param_name
+            _insert_in_params(
                 {
-                    "name": dependency.dependency.alias or dependency.param_name,
+                    "name": name.capitalize(),
                     "in": "header",
                     "description": dependency.dependency.description,
                     "required": not _is_optional(dependency.signature),
                 },
             )
         elif isinstance(dependency.dependency, Path):
-            route_info["parameters"].append(
+            _insert_in_params(
                 {
                     "name": dependency.dependency.alias or dependency.param_name,
                     "in": "path",
@@ -147,8 +155,9 @@ def _add_route_def(  # noqa: C901
                 },
             )
 
+    route_info["parameters"] = list(params.values())
     openapi_schema["paths"][route.resource.canonical].update(
-        {method.lower(): deep_update(route_info, extra_openapi)},
+        {method.lower(): always_merger.merge(route_info, extra_openapi)},
     )
 
 
@@ -265,7 +274,7 @@ async def event_handler(app: web.Application) -> None:
     return event_handler
 
 
-def extra_openapi(additional_schema: Dict[str, Any]) -> Callable[..., Any]:
+def extra_openapi(additional_schema: Dict[str, Any]) -> Callable[[_T], _T]:
     """
     Add extra openapi schema.
 
@@ -276,8 +285,46 @@ def extra_openapi(additional_schema: Dict[str, Any]) -> Callable[..., Any]:
     :return: same function with new attributes.
     """
 
-    def decorator(func: Any) -> Any:
-        func.__extra_openapi__ = additional_schema
+    def decorator(func: _T) -> _T:
+        func.__extra_openapi__ = additional_schema  # type: ignore
+        return func
+
+    return decorator
+
+
+def openapi_response(
+    status: int,
+    model: Any,
+    *,
+    content_type: str = "application/json",
+    description: Optional[str] = None,
+) -> Callable[[_T], _T]:
+    """
+    Add response schema to the endpoint.
+
+    This function takes a status and model,
+    which is going to represent the response.
+
+    :param status: Status of a response.
+    :param model: Response model.
+    :param content_type: Content-type of a response.
+    :param description: Response's description.
+
+    :returns: decorator that modifies your function.
+    """
+
+    def decorator(func: _T) -> _T:
+        openapi = getattr(func, "__extra_openapi__", {})
+        adapter: "pydantic.TypeAdapter[Any]" = pydantic.TypeAdapter(model)
+        responses = openapi.get("responses", {})
+        status_response = responses.get(status, {})
+        if not status_response:
+            status_response["description"] = description
+        status_response["content"] = status_response.get("content", {})
+        status_response["content"][content_type] = {"schema": adapter.json_schema()}
+        responses[status] = status_response
+        openapi["responses"] = responses
+        func.__extra_openapi__ = openapi  # type: ignore
         return func
 
     return decorator
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -28,6 +28,7 @@ python = "^3.8.1"
 aiohttp = "^3"
 taskiq-dependencies = "^1"
 pydantic = "^2"
+deepmerge = "^1.1.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.1.2"
diff --git a/tests/test_isoptional.py b/tests/test_isoptional.py
@@ -1,4 +1,5 @@
 import inspect
+import sys
 from typing import Optional, Union
 
 import pytest
@@ -42,7 +43,7 @@ def tfunc(param: Union[int, str]):
     assert not _is_optional(param)
 
 
-@pytest.mark.skip("We doesn't support 3.10 annotation style yet.")
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="Unsupported syntax")
 def test_new_type_style():
     def tfunc(param: "int | None"):
         """Nothing."""
@@ -52,7 +53,7 @@ def tfunc(param: "int | None"):
     assert _is_optional(param)
 
 
-@pytest.mark.skip("We doesn't support string annotations yet.")
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="Unsupported syntax")
 def test_string_annotation():
     def tfunc(param: "int | None"):
         """Nothing."""
diff --git a/tests/test_swagger.py b/tests/test_swagger.py
