You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

148 lines
5.0 KiB

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