schemas.py 5.1 KB

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