Merge pull request #11 from taskiq-python/bugfix/generics

s3rius · web-flow · commit 8959aa828045 · 2023-07-26T15:03:59.000+04:00
diff --git a/.flake8 b/.flake8
@@ -122,6 +122,10 @@ per-file-ignores =
   ; Found wrong metadata variable
   WPS410,
 
+  swagger.py:
+  ; Too many local variables
+  WPS210,
+
 exclude =
   ./.git,
   ./venv,
diff --git a/aiohttp_deps/swagger.py b/aiohttp_deps/swagger.py
@@ -13,6 +13,7 @@
 
 _T = TypeVar("_T")  # noqa: WPS111
 
+REF_TEMPLATE = "#/components/schemas/{model}"
 SCHEMA_KEY = "openapi_schema"
 SWAGGER_HTML_TEMPALTE = """
 <html lang="en">
@@ -75,12 +76,13 @@ def dummy(_var: annotation.annotation) -> None:  # type: ignore
     return var == Optional[var]
 
 
-def _add_route_def(  # noqa: C901, WPS210
+def _add_route_def(  # noqa: C901, WPS210, WPS211
     openapi_schema: Dict[str, Any],
     route: web.ResourceRoute,
     method: str,
     graph: DependencyGraph,
     extra_openapi: Dict[str, Any],
+    extra_openapi_schemas: Dict[str, Any],
 ) -> None:
     route_info: Dict[str, Any] = {
         "description": inspect.getdoc(graph.target),
@@ -90,6 +92,9 @@ def _add_route_def(  # noqa: C901, WPS210
     if route.resource is None:  # pragma: no cover
         return
 
+    if extra_openapi_schemas:
+        openapi_schema["components"]["schemas"].update(extra_openapi_schemas)
+
     params: Dict[tuple[str, str], Any] = {}
 
     def _insert_in_params(data: Dict[str, Any]) -> None:
@@ -114,9 +119,9 @@ def _insert_in_params(data: Dict[str, Any]) -> None:
             ):
                 input_schema = pydantic.TypeAdapter(
                     dependency.signature.annotation,
-                ).json_schema()
+                ).json_schema(ref_template=REF_TEMPLATE)
                 openapi_schema["components"]["schemas"].update(
-                    input_schema.pop("definitions", {}),
+                    input_schema.pop("$defs", {}),
                 )
                 route_info["requestBody"] = {
                     "content": {content_type: {"schema": input_schema}},
@@ -216,13 +221,19 @@ async def event_handler(app: web.Application) -> None:
                     "__extra_openapi__",
                     {},
                 )
+                extra_schemas = getattr(
+                    route._handler.original_handler,
+                    "__extra_openapi_schemas__",
+                    {},
+                )
                 try:
                     _add_route_def(
                         openapi_schema,
                         route,  # type: ignore
                         route.method,
                         route._handler.graph,
                         extra_openapi=extra_openapi,
+                        extra_openapi_schemas=extra_schemas,
                     )
                 except Exception as exc:  # pragma: no cover
                     logger.warn(
@@ -234,20 +245,23 @@ async def event_handler(app: web.Application) -> None:
             elif isinstance(route._handler, InjectableViewHandler):
                 for key, graph in route._handler.graph_map.items():
                     extra_openapi = getattr(
-                        getattr(
-                            route._handler.original_handler,
-                            key,
-                        ),
+                        getattr(route._handler.original_handler, key),
                         "__extra_openapi__",
                         {},
                     )
+                    extra_schemas = getattr(
+                        getattr(route._handler.original_handler, key),
+                        "__extra_openapi_schemas__",
+                        {},
+                    )
                     try:
                         _add_route_def(
                             openapi_schema,
                             route,  # type: ignore
                             key,
                             graph,
                             extra_openapi=extra_openapi,
+                            extra_openapi_schemas=extra_schemas,
                         )
                     except Exception as exc:  # pragma: no cover
                         logger.warn(
@@ -315,16 +329,20 @@ def openapi_response(
 
     def decorator(func: _T) -> _T:
         openapi = getattr(func, "__extra_openapi__", {})
+        openapi_schemas = getattr(func, "__extra_openapi_schemas__", {})
         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()}
+        response_schema = adapter.json_schema(ref_template=REF_TEMPLATE)
+        openapi_schemas.update(response_schema.pop("$defs", {}))
+        status_response["content"][content_type] = {"schema": response_schema}
         responses[status] = status_response
         openapi["responses"] = responses
         func.__extra_openapi__ = openapi  # type: ignore
+        func.__extra_openapi_schemas__ = openapi_schemas  # type: ignore
         return func
 
     return decorator
diff --git a/tests/test_swagger.py b/tests/test_swagger.py
@@ -1,5 +1,5 @@
 from collections import deque
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Generic, Optional, TypeVar
 
 import pytest
 from aiohttp import web
@@ -20,6 +20,21 @@
 from tests.conftest import ClientGenerator
 
 
+def follow_ref(ref: str, data: Dict[str, Any]) -> Dict[str, Any]:
+    """Function for following openapi references."""
+    components = deque(ref.split("/"))
+    current_model = None
+    while components:
+        component = components.popleft()
+        if component.strip() == "#":
+            current_model = data
+            continue
+        current_model = current_model.get(component)
+        if current_model is None:
+            return {}
+    return current_model
+
+
 def get_schema_by_ref(full_schema: Dict[str, Any], ref: str):
     ref_path = deque(ref.split("/"))
     current_schema = full_schema
@@ -141,19 +156,26 @@ async def my_handler(body=Depends(Json())):
     assert resp.status == 200
     resp_json = await resp.json()
     handler_info = resp_json["paths"]["/a"]["get"]
-    print(handler_info)
     assert handler_info["requestBody"]["content"]["application/json"] == {}
 
 
 @pytest.mark.anyio
-async def test_json_untyped(
+async def test_json_generic(
     my_app: web.Application,
     aiohttp_client: ClientGenerator,
 ):
     OPENAPI_URL = "/my_api_def.json"
     my_app.on_startup.append(setup_swagger(schema_url=OPENAPI_URL))
 
-    async def my_handler(body=Depends(Json())):
+    T = TypeVar("T")
+
+    class First(BaseModel):
+        name: str
+
+    class Second(BaseModel, Generic[T]):
+        data: T
+
+    async def my_handler(body: Second[First] = Depends(Json())):
         """Nothing."""
 
     my_app.router.add_get("/a", my_handler)
@@ -163,7 +185,10 @@ async def my_handler(body=Depends(Json())):
     assert resp.status == 200
     resp_json = await resp.json()
     handler_info = resp_json["paths"]["/a"]["get"]
-    assert {} == handler_info["requestBody"]["content"]["application/json"]
+    schema = handler_info["requestBody"]["content"]["application/json"]["schema"]
+    first_ref = schema["properties"]["data"]["$ref"]
+    first_obj = follow_ref(first_ref, resp_json)
+    assert "name" in first_obj["properties"]
 
 
 @pytest.mark.anyio
@@ -438,7 +463,6 @@ async def my_handler():
     resp_json = await resp.json()
 
     handler_info = resp_json["paths"]["/a"]["get"]
-    print(handler_info)
     assert handler_info["responses"] == {"200": {}}
 
 
@@ -495,7 +519,6 @@ async def my_handler(
     assert resp.status == 200
     resp_json = await resp.json()
     params = resp_json["paths"]["/a"]["get"]["parameters"]
-    print(params)
     assert len(params) == 1
     assert params[0]["name"] == "Head"
     assert params[0]["required"]
@@ -562,3 +585,34 @@ async def my_handler():
     assert "200" in route_info["responses"]
     assert "application/json" in route_info["responses"]["200"]["content"]
     assert "application/xml" in route_info["responses"]["200"]["content"]
+
+
+@pytest.mark.anyio
+async def test_custom_responses_generics(
+    my_app: web.Application,
+    aiohttp_client: ClientGenerator,
+) -> None:
+    OPENAPI_URL = "/my_api_def.json"
+    my_app.on_startup.append(setup_swagger(schema_url=OPENAPI_URL))
+
+    T = TypeVar("T")
+
+    class First(BaseModel):
+        name: str
+
+    class Second(BaseModel, Generic[T]):
+        data: T
+
+    @openapi_response(200, Second[First])
+    async def my_handler():
+        """Nothing."""
+
+    my_app.router.add_get("/a", my_handler)
+    client = await aiohttp_client(my_app)
+    response = await client.get(OPENAPI_URL)
+    resp_json = await response.json()
+    first_ref = resp_json["paths"]["/a"]["get"]["responses"]["200"]["content"][
+        "application/json"
+    ]["schema"]["properties"]["data"]["$ref"]
+    first_obj = follow_ref(first_ref, resp_json)
+    assert "name" in first_obj["properties"]
