Merge pull request #22 from koxudaxi/support_security

Support security

Author: koxudaxi

README.md

@@ -262,6 +262,7 @@ You can use below variables in jinja2 template
   - `operation.response` response object
   - `operation.function_name` function name is created `operationId` or `METHOD` + `Path` 
   - `operation.snake_case_arguments` Snake-cased function arguments
+  - `operation.security` [Security](https://swagger.io/docs/specification/authentication/)
 
 ### default template 
 `main.jinja2`

fastapi_code_generator/parser.py

@@ -89,6 +89,7 @@ class Operation(CachedPropertyModel):
     responses: Dict[UsefulStr, Any] = {}
     requestBody: Dict[str, Any] = {}
     imports: List[Import] = []
+    security: Optional[List[Dict[str, List[str]]]] = None
 
     @cached_property
     def root_path(self) -> UsefulStr:
@@ -299,6 +300,7 @@ class Operations(BaseModel):
     options: Optional[Operation] = None
     trace: Optional[Operation] = None
     path: UsefulStr
+    security: Optional[List[Dict[str, List[str]]]] = []
 
     @root_validator(pre=True)
     def inject_path_and_type_to_operation(cls, values: Dict[str, Any]) -> Any:
@@ -311,30 +313,40 @@ def inject_path_and_type_to_operation(cls, values: Dict[str, Any]) -> Any:
             },
             path=path,
             parameters=values.get('parameters', []),
+            security=values.get('security'),
         )
 
     @root_validator
-    def inject_parameters_to_operation(cls, values: Dict[str, Any]) -> Any:
-        if parameters := values.get('parameters'):
-            for operation_name in OPERATION_NAMES:
-                if operation := values.get(operation_name):
+    def inject_parameters_and_security_to_operation(cls, values: Dict[str, Any]) -> Any:
+        security = values.get('security')
+        for operation_name in OPERATION_NAMES:
+            if operation := values.get(operation_name):
+                if parameters := values.get('parameters'):
                     operation.parameters.extend(parameters)
+                if security is not None and operation.security is None:
+                    operation.security = security
+
         return values
 
 
 class Path(CachedPropertyModel):
     path: UsefulStr
     operations: Optional[Operations] = None
+    security: Optional[List[Dict[str, List[str]]]] = []
 
     @root_validator(pre=True)
     def validate_root(cls, values: Dict[str, Any]) -> Any:
         if path := values.get('path'):
             if isinstance(path, str):
                 if operations := values.get('operations'):
                     if isinstance(operations, dict):
+                        security = values.get('security', [])
                         return {
                             'path': path,
-                            'operations': dict(**operations, path=path),
+                            'operations': dict(
+                                **operations, path=path, security=security
+                            ),
+                            'security': security,
                         }
         return values
 
@@ -379,15 +391,21 @@ def __init__(
 
     def parse(self) -> ParsedObject:
         openapi = load_json_or_yaml(self.input_text)
-        return self.parse_paths(openapi["paths"])
+        return self.parse_paths(openapi)
+
+    def parse_security(
+        self, openapi: Dict[str, Any]
+    ) -> Optional[List[Dict[str, List[str]]]]:
+        return openapi.get('security')
 
-    def parse_paths(self, paths: Dict[str, Any]) -> ParsedObject:
+    def parse_paths(self, openapi: Dict[str, Any]) -> ParsedObject:
+        security = self.parse_security(openapi)
         return ParsedObject(
             [
                 operation
-                for path_name, operations in paths.items()
+                for path_name, operations in openapi['paths'].items()
                 for operation in Path(
-                    path=UsefulStr(path_name), operations=operations
+                    path=UsefulStr(path_name), operations=operations, security=security
                 ).exists_operations
             ]
         )

pyproject.toml

@@ -53,13 +53,15 @@ freezegun = "^0.3.15"
 line-length = 88
 skip-string-normalization = true
 target-version = ['py38']
+exclude = '(tests/data|\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|_build|buck-out|build|dist|.*\/models\.py.*|.*\/models\/.*)'
 
 [tool.isort]
 multi_line_output = 3
 include_trailing_comma = true
 force_grid_wrap = 0
 use_parentheses = true
 line_length = 88
+skip = "tests/data"
 
 [tool.pydantic-pycharm-plugin.parsable-types]
 # str field may parse int and float

tests/data/custom_template/security/main.jinja2

@@ -0,0 +1,47 @@
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from fastapi import Depends, FastAPI, HTTPException, Query
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+from starlette import status
+
+{{ imports }}
+
+app = FastAPI()
+
+
+DUMMY_CREDENTIALS = 'abcdefg'
+
+
+class User(BaseModel):
+    username: str
+    email: str
+
+
+def get_dummy_user(token: str) -> User:
+    return User(username=token, email='[email protected]')
+
+
+async def valid_token(auth: HTTPAuthorizationCredentials = Depends(HTTPBearer())) -> str:
+    if auth.credentials == DUMMY_CREDENTIALS:
+        return 'dummy'
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid authentication credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
+
+async def valid_current_user(token: str = Depends(valid_token)) -> User:
+    return get_dummy_user(token)
+
+
+
+{% for operation in operations %}
+@app.{{operation.type}}('{{operation.snake_case_path}}', response_model={{operation.response}})
+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}}:
+    pass
+{% endfor %}

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

@@ -0,0 +1,80 @@
+# generated by fastapi-codegen:
+#   filename:  custom_security.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from fastapi import Depends, FastAPI, HTTPException, Query
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from starlette import status
+
+from .models import Pet, PetForm
+
+app = FastAPI()
+
+
+DUMMY_CREDENTIALS = 'abcdefg'
+
+
+class User(BaseModel):
+    username: str
+    email: str
+
+
+def get_dummy_user(token: str) -> User:
+    return User(username=token, email='[email protected]')
+
+
+async def valid_token(
+    auth: HTTPAuthorizationCredentials = Depends(HTTPBearer()),
+) -> str:
+    if auth.credentials == DUMMY_CREDENTIALS:
+        return 'dummy'
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Invalid authentication credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
+
+async def valid_current_user(token: str = Depends(valid_token)) -> User:
+    return get_dummy_user(token)
+
+
+@app.get('/food/{food_id}', response_model=None)
+def show_food_by_id(food_id: str, user: User = Depends(valid_current_user)) -> None:
+    pass
+
+
+@app.get('/pets', response_model=List[Pet])
+def list_pets(
+    limit: Optional[int] = 0,
+    home_address: Optional[str] = Query('Unknown', alias='HomeAddress'),
+    kind: Optional[str] = 'dog',
+) -> List[Pet]:
+    pass
+
+
+@app.post('/pets', response_model=None)
+def post_pets(body: PetForm, user: User = Depends(valid_current_user)) -> None:
+    pass
+
+
+@app.get('/pets/{pet_id}', response_model=Pet)
+def show_pet_by_id(
+    pet_id: str = Query(..., alias='petId'), user: User = Depends(valid_current_user)
+) -> Pet:
+    pass
+
+
+@app.put('/pets/{pet_id}', response_model=None)
+def put_pets_pet_id(
+    pet_id: str = Query(..., alias='petId'),
+    body: PetForm = None,
+    user: User = Depends(valid_current_user),
+) -> None:
+    pass

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

@@ -0,0 +1,23 @@
+# generated by datamodel-codegen:
+#   filename:  custom_security.yaml
+#   timestamp: 2020-06-19T00:00:00+00:00
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class Pet(BaseModel):
+    id: int
+    name: str
+    tag: Optional[str] = None
+
+
+class Error(BaseModel):
+    code: int
+    message: str
+
+
+class PetForm(BaseModel):
+    name: Optional[str] = None
+    age: Optional[int] = None