feat: openapi response schema (#960)

* feat: openapi response schema * add tests * feat: handle robyn.Response * update import * docs: add * Update docs_src/src/pages/documentation/example_app/openapi.mdx --------- Co-authored-by: Sanskar Jethi <[email protected]>
sparckles · Sep 24, 2024 · 399f5b1 · 399f5b1
1 parent e6872af
commit 399f5b1
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 29 deletions.
diff --git a/docs_src/src/pages/documentation/api_reference/openapi.mdx b/docs_src/src/pages/documentation/api_reference/openapi.mdx
@@ -217,11 +217,13 @@ app.include_router(subrouter)
 
 ## Other Specification Params
 
-We support all the params mentioned in the latest OpenAPI specifications (https://swagger.io/specification/). See an example of using request body below.
+We support all the params mentioned in the latest OpenAPI specifications (https://swagger.io/specification/). See an example using request & response bodies below:
 
-<CodeGroup title="Request Body">
+<CodeGroup title="Request & Response Body">
 
 ```python {{ title: 'untyped' }}
+from robyn.types import JSONResponse
+
 class Initial(TypedDict):
     is_present: bool
     letter: Optional[str]
@@ -240,12 +242,19 @@ class CreateItemBody(TypedDict):
     tax: float
 
 
+class CreateResponse(JSONResponse):
+    success: bool
+    items_changed: int
+
+
 @app.post("/")
-def create_item(request, body=CreateItemBody):
-    return request.body
+def create_item(request: Request, body=CreateItemBody) -> CreateResponse:
+    return CreateResponse(success=True, items_changed=2)
 ```
 
 ```python {{ title: 'typed' }}
+from robyn.types import JSONResponse
+
 class Initial(TypedDict):
     is_present: bool
     letter: Optional[str]
@@ -263,10 +272,14 @@ class CreateItemBody(TypedDict):
     price: float
     tax: float
 
+class CreateResponse(JSONResponse):
+    success: bool
+    items_changed: int
+
 
 @app.post("/")
-def create_item(request: Request, body=CreateItemBody):
-    return request.body
+def create_item(request: Request, body=CreateItemBody) -> CreateResponse:
+    return CreateResponse(success=True, items_changed=2)
 ```
 
 </CodeGroup>

diff --git a/docs_src/src/pages/documentation/example_app/openapi.mdx b/docs_src/src/pages/documentation/example_app/openapi.mdx
@@ -175,11 +175,13 @@ app.include_router(subrouter)
 
 ## Other Specification Params
 
-We support all the params mentioned in the latest OpenAPI specifications (https://swagger.io/specification/). See an example of using request body below.
+We support all the params mentioned in the latest OpenAPI specifications (https://swagger.io/specification/). See an example using request & response bodies below:
 
-<CodeGroup title="Request Body">
+<CodeGroup title="Request & Response Body">
 
 ```python {{ title: 'untyped' }}
+from robyn.types import JSONResponse
+
 class Initial(TypedDict):
     is_present: bool
     letter: Optional[str]
@@ -198,12 +200,19 @@ class CreateItemBody(TypedDict):
     tax: float
 
 
+class CreateResponse(JSONResponse):
+    success: bool
+    items_changed: int
+
+
 @app.post("/")
-def create_item(request, body=CreateItemBody):
-    return request.body
+def create_item(request: Request, body=CreateItemBody) -> CreateResponse:
+    return {"success": True, "items_changed": 2}
 ```
 
 ```python {{ title: 'typed' }}
+from robyn.types import JSONResponse
+
 class Initial(TypedDict):
     is_present: bool
     letter: Optional[str]
@@ -221,10 +230,14 @@ class CreateItemBody(TypedDict):
     price: float
     tax: float
 
+class CreateResponse(JSONResponse):
+    success: bool
+    items_changed: int
+
 
 @app.post("/")
-def create_item(request: Request, body=CreateItemBody):
-    return request.body
+def create_item(request: Request, body=CreateItemBody) -> CreateResponse:
+    return CreateResponse(success=True, items_changed=2)
 ```
 
 </CodeGroup>

diff --git a/integration_tests/base_routes.py b/integration_tests/base_routes.py
@@ -858,7 +858,7 @@ class CreateItemBody(TypedDict):
 
 
 @app.post("/openapi_request_body")
-def create_item(request, body=CreateItemBody):
+def create_item(request, body=CreateItemBody) -> CreateItemBody:
     return request.body
 
 

diff --git a/integration_tests/test_openapi.py b/integration_tests/test_openapi.py
@@ -114,3 +114,78 @@ def test_openapi_request_body():
     assert {"type": "null"} in openapi_spec["paths"][endpoint][route_type]["requestBody"]["content"]["application/json"]["schema"]["properties"]["name"][
         "properties"
     ]["initial"]["properties"]["letter"]["anyOf"]
+
+
+@pytest.mark.benchmark
+def test_openapi_response_body():
+    openapi_spec = get("/openapi.json").json()
+
+    assert isinstance(openapi_spec, dict)
+
+    route_type = "post"
+    endpoint = "/openapi_request_body"
+
+    assert endpoint in openapi_spec["paths"]
+    assert route_type in openapi_spec["paths"][endpoint]
+    assert "responses" in openapi_spec["paths"][endpoint][route_type]
+    assert "200" in openapi_spec["paths"][endpoint][route_type]["responses"]
+
+    assert openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["description"] == "Successful Response"
+
+    assert "content" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]
+
+    assert "application/json" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]
+    assert "schema" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]
+    assert "properties" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]
+
+    assert "name" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]
+    assert "description" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]
+    assert "price" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]
+    assert "tax" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]
+
+    assert (
+        "string"
+        == openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["description"]["type"]
+    )
+    assert "number" == openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["price"]["type"]
+    assert "number" == openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["tax"]["type"]
+
+    assert "object" == openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["type"]
+
+    assert (
+        "first" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"]
+    )
+    assert (
+        "second" in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"]
+    )
+    assert (
+        "initial"
+        in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"]
+    )
+
+    assert (
+        "object"
+        in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"][
+            "initial"
+        ]["type"]
+    )
+
+    assert (
+        "is_present"
+        in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"][
+            "initial"
+        ]["properties"]
+    )
+    assert (
+        "letter"
+        in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"]["properties"][
+            "initial"
+        ]["properties"]
+    )
+
+    assert {"type": "string"} in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"][
+        "properties"
+    ]["initial"]["properties"]["letter"]["anyOf"]
+    assert {"type": "null"} in openapi_spec["paths"][endpoint][route_type]["responses"]["200"]["content"]["application/json"]["schema"]["properties"]["name"][
+        "properties"
+    ]["initial"]["properties"]["letter"]["anyOf"]
diff --git a/robyn/openapi.py b/robyn/openapi.py
@@ -4,6 +4,7 @@
 from inspect import Signature
 from typing import Callable, Dict, List, Optional, TypedDict, Any
 
+from robyn import Response
 from robyn.responses import FileResponse, html
 
 
@@ -162,21 +163,24 @@ def add_openapi_path_obj(self, route_type: str, endpoint: str, openapi_name: str
 
         query_params = None
         request_body = None
+        return_annotation = None
 
         signature = inspect.signature(handler)
         openapi_description = inspect.getdoc(handler) or ""
 
-        if signature and "query_params" in signature.parameters:
-            query_params = signature.parameters["query_params"].default
+        if signature:
+            if "query_params" in signature.parameters:
+                query_params = signature.parameters["query_params"].default
 
-        if signature and "body" in signature.parameters:
-            request_body = signature.parameters["body"].default
+            if "body" in signature.parameters:
+                request_body = signature.parameters["body"].default
 
-        return_annotation = signature.return_annotation
+            if signature.return_annotation is not Signature.empty:
+                return_annotation = signature.return_annotation
 
-        return_type = "text/plain" if return_annotation == Signature.empty or return_annotation is str else "application/json"
-
-        modified_endpoint, path_obj = self.get_path_obj(endpoint, openapi_name, openapi_description, openapi_tags, query_params, request_body, return_type)
+        modified_endpoint, path_obj = self.get_path_obj(
+            endpoint, openapi_name, openapi_description, openapi_tags, query_params, request_body, return_annotation
+        )
 
         if modified_endpoint not in self.openapi_spec["paths"]:
             self.openapi_spec["paths"][modified_endpoint] = {}
@@ -201,7 +205,7 @@ def get_path_obj(
         tags: List[str],
         query_params: Optional[TypedDict],
         request_body: Optional[TypedDict],
-        return_type: str,
+        return_annotation: Optional[TypedDict],
     ) -> (str, dict):
         """
         Get the "path" openapi object according to spec
@@ -212,7 +216,7 @@ def get_path_obj(
         @param tags: List[str] for grouping of endpoints
         @param query_params: Optional[TypedDict] query params for the function
         @param request_body: Optional[TypedDict] request body for the function
-        @param return_type: str return type of the endpoint handler
+        @param return_annotation: Optional[TypedDict] return type of the endpoint handler
 
         @return: (str, dict) a tuple containing the endpoint with path params wrapped in braces and the "path" openapi object
         according to spec
@@ -226,7 +230,6 @@ def get_path_obj(
             "description": description,
             "parameters": [],
             "tags": tags,
-            "responses": {"200": {"description": "Successful Response", "content": {return_type: {"schema": {}}}}},
         }
 
         # robyn has paths like /:url/:etc whereas openapi requires path like /{url}/{path}
@@ -273,7 +276,7 @@ def get_path_obj(
             properties = {}
 
             for body_item in request_body.__annotations__:
-                properties[body_item] = self.get_properties_object(body_item, request_body.__annotations__[body_item])
+                properties[body_item] = self.get_schema_object(body_item, request_body.__annotations__[body_item])
 
             request_body_object = {
                 "content": {
@@ -288,6 +291,15 @@ def get_path_obj(
 
             openapi_path_object["requestBody"] = request_body_object
 
+        response_schema = {}
+        response_type = "text/plain"
+
+        if return_annotation and return_annotation is not Response:
+            response_type = "application/json"
+            response_schema = self.get_schema_object("response object", return_annotation)
+
+        openapi_path_object["responses"] = {"200": {"description": "Successful Response", "content": {response_type: {"schema": response_schema}}}}
+
         return endpoint_with_path_params_wrapped_in_braces, openapi_path_object
 
     def get_openapi_type(self, typed_dict: TypedDict) -> str:
@@ -313,9 +325,9 @@ def get_openapi_type(self, typed_dict: TypedDict) -> str:
         # default to "string" if type is not found
         return "string"
 
-    def get_properties_object(self, parameter: str, param_type: Any) -> dict:
+    def get_schema_object(self, parameter: str, param_type: Any) -> dict:
         """
-        Get the properties object for request body
+        Get the schema object for request/response body
 
         @param parameter: name of the parameter
         @param param_type: Any the type to be inferred
@@ -351,7 +363,7 @@ def get_properties_object(self, parameter: str, param_type: Any) -> dict:
             properties["properties"] = {}
 
             for e in param_type.__annotations__:
-                properties["properties"][e] = self.get_properties_object(e, param_type.__annotations__[e])
+                properties["properties"][e] = self.get_schema_object(e, param_type.__annotations__[e])
 
         properties["type"] = "object"
 

diff --git a/robyn/types.py b/robyn/types.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, TypedDict
 
 
 @dataclass
@@ -16,3 +16,7 @@ def as_list(self):
             self.show_files_listing,
             self.index_file,
         ]
+
+
+class JSONResponse(TypedDict):
+    pass