schemas.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148
  1. from __future__ import annotations
  2. import inspect
  3. import re
  4. from collections.abc import Callable
  5. from typing import Any, NamedTuple
  6. from starlette.requests import Request
  7. from starlette.responses import Response
  8. from starlette.routing import BaseRoute, Host, Mount, Route
  9. try:
  10. import yaml
  11. except ModuleNotFoundError: # pragma: no cover
  12. yaml = None # type: ignore[assignment]
  13. class OpenAPIResponse(Response):
  14. media_type = "application/vnd.oai.openapi"
  15. def render(self, content: Any) -> bytes:
  16. assert yaml is not None, "`pyyaml` must be installed to use OpenAPIResponse."
  17. assert isinstance(content, dict), "The schema passed to OpenAPIResponse should be a dictionary."
  18. return yaml.dump(content, default_flow_style=False).encode("utf-8")
  19. class EndpointInfo(NamedTuple):
  20. path: str
  21. http_method: str
  22. func: Callable[..., Any]
  23. _remove_converter_pattern = re.compile(r":\w+}")
  24. class BaseSchemaGenerator:
  25. def get_schema(self, routes: list[BaseRoute]) -> dict[str, Any]:
  26. raise NotImplementedError() # pragma: no cover
  27. def get_endpoints(self, routes: list[BaseRoute]) -> list[EndpointInfo]:
  28. """
  29. Given the routes, yields the following information:
  30. - path
  31. eg: /users/
  32. - http_method
  33. one of 'get', 'post', 'put', 'patch', 'delete', 'options'
  34. - func
  35. method ready to extract the docstring
  36. """
  37. endpoints_info: list[EndpointInfo] = []
  38. for route in routes:
  39. if isinstance(route, Mount | Host):
  40. routes = route.routes or []
  41. if isinstance(route, Mount):
  42. path = self._remove_converter(route.path)
  43. else:
  44. path = ""
  45. sub_endpoints = [
  46. EndpointInfo(
  47. path="".join((path, sub_endpoint.path)),
  48. http_method=sub_endpoint.http_method,
  49. func=sub_endpoint.func,
  50. )
  51. for sub_endpoint in self.get_endpoints(routes)
  52. ]
  53. endpoints_info.extend(sub_endpoints)
  54. elif not isinstance(route, Route) or not route.include_in_schema:
  55. continue
  56. elif inspect.isfunction(route.endpoint) or inspect.ismethod(route.endpoint):
  57. path = self._remove_converter(route.path)
  58. for method in route.methods or ["GET"]:
  59. if method == "HEAD":
  60. continue
  61. endpoints_info.append(EndpointInfo(path, method.lower(), route.endpoint))
  62. else:
  63. path = self._remove_converter(route.path)
  64. for method in ["get", "post", "put", "patch", "delete", "options"]:
  65. if not hasattr(route.endpoint, method):
  66. continue
  67. func = getattr(route.endpoint, method)
  68. endpoints_info.append(EndpointInfo(path, method.lower(), func))
  69. return endpoints_info
  70. def _remove_converter(self, path: str) -> str:
  71. """
  72. Remove the converter from the path.
  73. For example, a route like this:
  74. Route("/users/{id:int}", endpoint=get_user, methods=["GET"])
  75. Should be represented as `/users/{id}` in the OpenAPI schema.
  76. """
  77. return _remove_converter_pattern.sub("}", path)
  78. def parse_docstring(self, func_or_method: Callable[..., Any]) -> dict[str, Any]:
  79. """
  80. Given a function, parse the docstring as YAML and return a dictionary of info.
  81. """
  82. docstring = func_or_method.__doc__
  83. if not docstring:
  84. return {}
  85. assert yaml is not None, "`pyyaml` must be installed to use parse_docstring."
  86. # We support having regular docstrings before the schema
  87. # definition. Here we return just the schema part from
  88. # the docstring.
  89. docstring = docstring.split("---")[-1]
  90. parsed = yaml.safe_load(docstring)
  91. if not isinstance(parsed, dict):
  92. # A regular docstring (not yaml formatted) can return
  93. # a simple string here, which wouldn't follow the schema.
  94. return {}
  95. return parsed
  96. def OpenAPIResponse(self, request: Request) -> Response:
  97. routes = request.app.routes
  98. schema = self.get_schema(routes=routes)
  99. return OpenAPIResponse(schema)
  100. class SchemaGenerator(BaseSchemaGenerator):
  101. def __init__(self, base_schema: dict[str, Any]) -> None:
  102. self.base_schema = base_schema
  103. def get_schema(self, routes: list[BaseRoute]) -> dict[str, Any]:
  104. schema = dict(self.base_schema)
  105. schema.setdefault("paths", {})
  106. endpoints_info = self.get_endpoints(routes)
  107. for endpoint in endpoints_info:
  108. parsed = self.parse_docstring(endpoint.func)
  109. if not parsed:
  110. continue
  111. if endpoint.path not in schema["paths"]:
  112. schema["paths"][endpoint.path] = {}
  113. schema["paths"][endpoint.path][endpoint.http_method] = parsed
  114. return schema