Repository: vitalik/django-ninja
Pull request #37 (2020-11-20 13:08:53)

Implement multiple response schemas
This implements multiple response schemas using the dict approach discussed in issue #16.


Patch:
diff --git a/ninja/openapi/schema.py b/ninja/openapi/schema.py
index 28d1794..58710ce 100644
--- a/ninja/openapi/schema.py
+++ b/ninja/openapi/schema.py
@@ -114,14 +114,26 @@ class OpenAPISchema(OrderedDict):
 
     def responses(self, operation):
         if operation.response_model:
-            schema, _ = self._create_schema_from_model(operation.response_model)
-
-            return {
-                200: {
-                    "description": "OK",
-                    "content": {"application/json": {"schema": schema}},
+            if not isinstance(operation.response_model, dict):
+                schema, _ = self._create_schema_from_model(operation.response_model)
+                return {
+                    200: {
+                        "description": "OK",
+                        "content": {"application/json": {"schema": schema}},
+                    }
                 }
-            }
+            
+            responses = {}
+            for status, model in operation.response_model.items():
+                schema, _ = self._create_schema_from_model(model)
+                responses.update({
+                    status: {
+                        "description": "OK" if status > 100 and status < 300 else "Error",
+                        "content": {"application/json": {"schema": schema}},
+                    }
+                })
+
+            return responses
         else:
             return {200: {"description": "OK"}}
 
diff --git a/ninja/operation.py b/ninja/operation.py
index f52615e..0ef3eca 100644
--- a/ninja/operation.py
+++ b/ninja/operation.py
@@ -3,7 +3,7 @@ import django
 from django.http import HttpResponse, HttpResponseNotAllowed
 from typing import Callable, List, Any, Union, Optional, Sequence
 from ninja.responses import Response
-from ninja.errors import InvalidInput
+from ninja.errors import InvalidInput, ConfigError
 from ninja.constants import NOT_SET
 from ninja.schema import Schema
 from ninja.signature import ViewSignature, is_async
@@ -32,7 +32,11 @@ class Operation:
 
         self.signature = ViewSignature(self.path, self.view_func)
         self.models = self.signature.models
-        self.response_model = self._create_response_model(response)
+
+        if isinstance(response, dict):
+            self.response_model = self._create_response_model_multiple(response)
+        else:
+            self.response_model = self._create_response_model(response)
 
     def run(self, request, **kw):
         error = self._run_checks(request)
@@ -82,11 +86,21 @@ class Operation:
             return result
         if self.response_model is None:
             return Response(result)
+        
+        status = 200
+        response_model = self.response_model
+        if isinstance(result, tuple) and len(result) == 2:
+            status = result[0]
+            result = result[1]
+        if isinstance(response_model, dict):
+            if not status in response_model.keys():
+                raise ConfigError(f'Schema for status {status} is not set in response')
+            response_model = response_model[status]
 
         resp_object = ResponseObject(result)
         # ^ we need object because getter_dict seems work only with from_orm
-        result = self.response_model.from_orm(resp_object).dict()["response"]
-        return Response(result)
+        result = response_model.from_orm(resp_object).dict()["response"]
+        return Response(result, status=status)
 
     def _get_values(self, request, path_params):
         values, errors = {}, []
@@ -102,6 +116,11 @@ class Operation:
                 errors.extend(items)
         return values, errors
 
+    def _create_response_model_multiple(self, response_param):
+        for status, model in response_param.items():
+            response_param[status] = self._create_response_model(model)
+        return response_param
+
     def _create_response_model(self, response_param):
         if response_param is None:
             return


Test patch:
diff --git a/tests/test_response_multiple.py b/tests/test_response_multiple.py
new file mode 100644
index 0000000..9cb062a
--- /dev/null
+++ b/tests/test_response_multiple.py
@@ -0,0 +1,105 @@
+from ninja import responses
+from ninja.errors import ConfigError
+import pytest
+from pydantic import ValidationError, BaseModel
+from ninja import Router
+from client import NinjaClient
+from typing import List, Union
+
+
+router = Router()
+
+
+# TODO: Return annotaion is not good
+# results are not always matching the type
+# and long functions will be even longer/unreadable...
+
+
+@router.get("/check_int", response={200: int})
+def check_int(request):
+    return 200, "1"
+
+
+@router.get("/check_int2", response={200: int})
+def check_int2(request):
+    return 200, "str"
+
+@router.get("/check_single_with_status", response=int)
+def check_single_with_status(request):
+    return 302, 1
+
+@router.get("/check_response_schema", response={400: int})
+def check_response_schema(request):
+    return 200, 1
+
+
+class User:
+    def __init__(self, id, name, password):
+        self.id = id
+        self.name = name
+        self.password = password
+
+
+class UserModel(BaseModel):
+    id: int
+    name: str
+    # skipping password output to responses
+
+    class Config:
+        orm_mode = True
+
+class ErrorModel(BaseModel):
+    detail: str
+
+
+@router.get("/check_model", response={200: UserModel, 202: UserModel})
+def check_model(request):
+    return 202, User(1, "John", "Password")
+
+
+@router.get("/check_list_model", response={200: List[UserModel]})
+def check_list_model(request):
+    return 200, [User(1, "John", "Password")]
+
+
+@router.get("/check_union", response={200: Union[int, UserModel], 400: ErrorModel})
+def check_union(request, q: int):
+    if q == 0:
+        return 200, 1
+    if q == 1:
+        return 200, User(1, "John", "Password")
+    if q == 2:
+        return 400, {"detail": "error"}
+    return "invalid"
+
+
+client = NinjaClient(router)
+
+
+@pytest.mark.parametrize(
+    "path,expected_status,expected_response",
+    [
+        ("/check_int", 200, 1),
+        ("/check_single_with_status", 302, 1),
+        ("/check_model", 202, {"id": 1, "name": "John"}),  # the password is skipped
+        ("/check_list_model", 200, [{"id": 1, "name": "John"}]),  # the password is skipped
+        ("/check_union?q=0", 200, 1),
+        ("/check_union?q=1", 200, {"id": 1, "name": "John"}),
+        ("/check_union?q=2", 400, {"detail": "error"}),
+    ],
+)
+def test_responses(path, expected_status, expected_response):
+    response = client.get(path)
+    assert response.status_code == expected_status, response.content
+    assert response.json() == expected_response
+
+
+def test_validates():
+    with pytest.raises(ValidationError):
+        client.get("/check_int2")
+
+    with pytest.raises(ValidationError):
+        client.get("/check_union?q=3")
+
+    with pytest.raises(ConfigError):
+        client.get("/check_response_schema")