Merge pull request #203 from rominf/rominf/add_additional_responses

Add additional responses

fastapi_code_generator/parser.py

@@ -121,6 +121,8 @@ class Operation(CachedPropertyModel):
     snake_case_arguments: str = ''
     request: Optional[Argument] = None
     response: str = ''
+    additional_responses: Dict[str, Dict[str, str]] = {}
+    return_type: str = ''
 
     @cached_property
     def type(self) -> UsefulStr:
@@ -380,11 +382,28 @@ def parse_responses(
             data_type = list(status_code_200.values())[0]
             if data_type:
                 self.data_types.append(data_type)
-            type_hint = data_type.type_hint  # TODO: change to lazy loading
         else:
-            type_hint = 'None'
+            data_type = DataType(type='None')
+        type_hint = data_type.type_hint  # TODO: change to lazy loading
         self._temporary_operation['response'] = type_hint
-
+        return_types = {type_hint: data_type}
+        for status_code, additional_responses in data_types.items():
+            if status_code != '200' and additional_responses:  # 200 is processed above
+                data_type = list(additional_responses.values())[0]
+                if data_type:
+                    self.data_types.append(data_type)
+                type_hint = data_type.type_hint  # TODO: change to lazy loading
+                self._temporary_operation.setdefault('additional_responses', {})[
+                    status_code
+                ] = {'model': type_hint}
+                return_types[type_hint] = data_type
+        if len(return_types) == 1:
+            return_type = next(iter(return_types.values()))
+        else:
+            return_type = DataType(data_types=list(return_types.values()))
+        if return_type:
+            self.data_types.append(return_type)
+        self._temporary_operation['return_type'] = return_type.type_hint
         return data_types
 
     def parse_operation(self, raw_operation: Dict[str, Any], path: List[str],) -> None:

fastapi_code_generator/template/main.jinja2

@@ -15,8 +15,19 @@ app = FastAPI(
 
 
 {% for operation in operations %}
-@app.{{operation.type}}('{{operation.snake_case_path}}', response_model={{operation.response}})
-def {{operation.function_name}}({{operation.snake_case_arguments}}) -> {{operation.response}}:
+@app.{{operation.type}}('{{operation.snake_case_path}}', response_model={{operation.response}}
+    {% if operation.additional_responses %}
+        , responses={
+            {% for status_code, models in operation.additional_responses.items() %}
+                '{{ status_code }}': {
+                {% for key, model in models.items() %}
+                    '{{ key }}': {{ model }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+                }{% if not loop.last %},{% endif %}
+            {% endfor %}
+        }
+    {% endif %})
+def {{operation.function_name}}({{operation.snake_case_arguments}}) -> {{operation.return_type}}:
     {%- if operation.summary %}
     """
     {{ operation.summary }}

tests/data/custom_template/security/main.jinja2

@@ -41,7 +41,19 @@ async def valid_current_user(token: str = Depends(valid_token)) -> User:
 
 
 {% for operation in operations %}
-@app.{{operation.type}}('{{operation.snake_case_path}}', response_model={{operation.response}}{% if operation.tags -%}, tags={{operation.tags}}{%- endif -%})
-def {{operation.function_name}}({{operation.snake_case_arguments}}{%- if operation.security -%}{%- if operation.snake_case_arguments -%}, {%- endif -%}user: User = Depends(valid_current_user){%- endif -%}) -> {{operation.response}}:
+@app.{{operation.type}}('{{operation.snake_case_path}}', response_model={{operation.response}}
+    {% if operation.tags -%}, tags={{operation.tags}}{%- endif -%}
+    {% if operation.additional_responses %}
+        , responses={
+        {% for status_code, models in operation.additional_responses.items() %}
+            '{{ status_code }}': {
+            {% for key, model in models.items() %}
+                '{{ key }}': {{ model }}{% if not loop.last %},{% endif %}
+            {% endfor %}
+            }{% if not loop.last %},{% endif %}
+        {% endfor %}
+        }
+    {% endif %})
+def {{operation.function_name}}({{operation.snake_case_arguments}}{%- if operation.security -%}{%- if operation.snake_case_arguments -%}, {%- endif -%}user: User = Depends(valid_current_user){%- endif -%}) -> {{operation.return_type}}:
     pass
 {% endfor %}

tests/data/expected/openapi/custom_template_security/custom_security/main.py

@@ -4,15 +4,15 @@
 
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from pydantic import BaseModel
 
 from fastapi import Depends, FastAPI, HTTPException, Path, Query
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 from starlette import status
 
-from .models import Pet, PetForm
+from .models import Error, Pet, PetForm
 
 app = FastAPI()
 
@@ -50,31 +50,45 @@ def show_food_by_id(food_id: str, user: User = Depends(valid_current_user)) -> N
     pass
 
 
-@app.get('/pets', response_model=List[Pet])
+@app.get('/pets', response_model=List[Pet], responses={'default': {'model': Error}})
 def list_pets(
     limit: Optional[int] = 0,
     home_address: Optional[str] = Query('Unknown', alias='HomeAddress'),
     kind: Optional[str] = 'dog',
-) -> List[Pet]:
+) -> Union[List[Pet], Error]:
     pass
 
 
-@app.post('/pets', response_model=None, tags=['pets'])
-def post_pets(body: PetForm, user: User = Depends(valid_current_user)) -> None:
+@app.post(
+    '/pets', response_model=None, tags=['pets'], responses={'default': {'model': Error}}
+)
+def post_pets(
+    body: PetForm, user: User = Depends(valid_current_user)
+) -> Union[None, Error]:
     pass
 
 
-@app.get('/pets/{pet_id}', response_model=Pet, tags=['pets'])
+@app.get(
+    '/pets/{pet_id}',
+    response_model=Pet,
+    tags=['pets'],
+    responses={'default': {'model': Error}},
+)
 def show_pet_by_id(
     pet_id: str = Path(..., alias='petId'), user: User = Depends(valid_current_user)
-) -> Pet:
+) -> Union[Pet, Error]:
     pass
 
 
-@app.put('/pets/{pet_id}', response_model=None, tags=['pets'])
+@app.put(
+    '/pets/{pet_id}',
+    response_model=None,
+    tags=['pets'],
+    responses={'default': {'model': Error}},
+)
 def put_pets_pet_id(
     pet_id: str = Path(..., alias='petId'),
     body: PetForm = None,
     user: User = Depends(valid_current_user),
-) -> None:
+) -> Union[None, Error]:
     pass

tests/data/expected/openapi/default_template/body_and_parameters/main.py

@@ -4,12 +4,13 @@
 
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from fastapi import FastAPI, Path, Query
 from starlette.requests import Request
 
 from .models import (
+    Error,
     Pet,
     PetForm,
     UserGetResponse,
@@ -39,8 +40,8 @@ def get_foo(foo: Optional[str] = None) -> str:
     pass
 
 
-@app.post('/food', response_model=None)
-def post_food(body: str) -> None:
+@app.post('/food', response_model=None, responses={'default': {'model': str}})
+def post_food(body: str) -> Union[None, str]:
     """
     Create a food
     """
@@ -57,38 +58,38 @@ def show_food_by_id(
     pass
 
 
-@app.get('/pets', response_model=List[Pet])
+@app.get('/pets', response_model=List[Pet], responses={'default': {'model': Error}})
 def list_pets(
     limit: Optional[int] = 0,
     home_address: Optional[str] = Query('Unknown', alias='HomeAddress'),
     kind: Optional[str] = 'dog',
-) -> List[Pet]:
+) -> Union[List[Pet], Error]:
     """
     List all pets
     """
     pass
 
 
-@app.post('/pets', response_model=None)
-def post_pets(body: PetForm) -> None:
+@app.post('/pets', response_model=None, responses={'default': {'model': Error}})
+def post_pets(body: PetForm) -> Union[None, Error]:
     """
     Create a pet
     """
     pass
 
 
-@app.get('/pets/{pet_id}', response_model=Pet)
-def show_pet_by_id(pet_id: str = Path(..., alias='petId')) -> Pet:
+@app.get('/pets/{pet_id}', response_model=Pet, responses={'default': {'model': Error}})
+def show_pet_by_id(pet_id: str = Path(..., alias='petId')) -> Union[Pet, Error]:
     """
     Info for a specific pet
     """
     pass
 
 
-@app.put('/pets/{pet_id}', response_model=None)
+@app.put('/pets/{pet_id}', response_model=None, responses={'default': {'model': Error}})
 def put_pets_pet_id(
     pet_id: str = Path(..., alias='petId'), body: PetForm = None
-) -> None:
+) -> Union[None, Error]:
     """
     update a pet
     """

tests/data/expected/openapi/default_template/oas_response_reference/main.py

@@ -4,11 +4,11 @@
 
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from fastapi import FastAPI
 
-from .models import Pet
+from .models import Error, Pet
 
 app = FastAPI(
     version='1.0.0',
@@ -18,8 +18,8 @@
 )
 
 
-@app.get('/pets', response_model=List[Pet])
-def list_pets(limit: Optional[int] = 0) -> List[Pet]:
+@app.get('/pets', response_model=List[Pet], responses={'default': {'model': Error}})
+def list_pets(limit: Optional[int] = 0) -> Union[List[Pet], Error]:
     """
     List all pets
     """

tests/data/expected/openapi/default_template/same_response_model_for_different_status_codes/main.py

@@ -0,0 +1,46 @@
+# generated by fastapi-codegen:
+#   filename:  same_response_model_for_different_status_codes.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Optional, Union
+
+from fastapi import FastAPI, Path
+
+from .models import Error, Pets
+
+app = FastAPI(
+    version='1.0.0',
+    title='Swagger Petstore',
+    license={'name': 'MIT'},
+    description='This description is for testing\nmulti-line\ndescription\n',
+)
+
+
+@app.get('/pets', response_model=Pets, responses={'default': {'model': Error}})
+def list_pets(limit: Optional[int] = None) -> Union[Pets, Error]:
+    """
+    List all pets
+    """
+    pass
+
+
+@app.post('/pets', response_model=None, responses={'default': {'model': Error}})
+def create_pets() -> Union[None, Error]:
+    """
+    Create a pet
+    """
+    pass
+
+
+@app.get(
+    '/pets/{pet_id}',
+    response_model=Pets,
+    responses={'404': {'model': Error}, 'default': {'model': Error}},
+)
+def show_pet_by_id(pet_id: str = Path(..., alias='petId')) -> Union[Pets, Error]:
+    """
+    Info for a specific pet
+    """
+    pass

tests/data/expected/openapi/default_template/same_response_model_for_different_status_codes/models.py

@@ -0,0 +1,24 @@
+# generated by fastapi-codegen:
+#   filename:  same_response_model_for_different_status_codes.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Pet(BaseModel):
+    id: int
+    name: str
+    tag: Optional[str] = None
+
+
+class Pets(BaseModel):
+    __root__: List[Pet] = Field(..., description='list of pet')
+
+
+class Error(BaseModel):
+    code: int
+    message: str

tests/data/expected/openapi/default_template/simple/main.py

@@ -4,11 +4,11 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Union
 
 from fastapi import FastAPI, Path
 
-from .models import Pets
+from .models import Error, Pets
 
 app = FastAPI(
     version='1.0.0',
@@ -18,24 +18,24 @@
 )
 
 
-@app.get('/pets', response_model=Pets)
-def list_pets(limit: Optional[int] = None) -> Pets:
+@app.get('/pets', response_model=Pets, responses={'default': {'model': Error}})
+def list_pets(limit: Optional[int] = None) -> Union[Pets, Error]:
     """
     List all pets
     """
     pass
 
 
-@app.post('/pets', response_model=None)
-def create_pets() -> None:
+@app.post('/pets', response_model=None, responses={'default': {'model': Error}})
+def create_pets() -> Union[None, Error]:
     """
     Create a pet
     """
     pass
 
 
-@app.get('/pets/{pet_id}', response_model=Pets)
-def show_pet_by_id(pet_id: str = Path(..., alias='petId')) -> Pets:
+@app.get('/pets/{pet_id}', response_model=Pets, responses={'default': {'model': Error}})
+def show_pet_by_id(pet_id: str = Path(..., alias='petId')) -> Union[Pets, Error]:
     """
     Info for a specific pet
     """