diff --git a/src/utils/templates.py b/src/utils/templates.py new file mode 100644 index 0000000..d0204da --- /dev/null +++ b/src/utils/templates.py @@ -0,0 +1,190 @@ +import http.server +import json +import shutil +import socketserver +import tempfile +from pathlib import Path +from typing import Any + +import click +import jinja2 +from jinja2 import BaseLoader + +from src.core.parser import parse_openapi_spec +from src.utils.examples import ExampleGenerator + + +class Jinja2Loader(BaseLoader): + def __init__(self, templates_dir: Path): + self.templates_dir = templates_dir + + def get_source(self, environment: jinja2.Environment, template: str) -> tuple: + path = self.templates_dir / template + if not path.exists(): + raise jinja2.TemplateNotFound(template) + return path.read_text(), str(path), lambda: True + + +def generate_html(spec_path: str, output_path: str, template_path: str | None = None) -> None: + spec = parse_openapi_spec(spec_path) + if template_path: + template_dir = Path(template_path).parent + else: + template_dir = Path(__file__).parent.parent / "templates" + loader = Jinja2Loader(template_dir) + env = jinja2.Environment(loader=loader) + env.filters["tojson"] = lambda x: json.dumps(x, indent=2) + template = env.get_template(Path(template_path).name if template_path else "html_template.html") + spec_dict = spec.model_dump() + components_schemas = spec_dict.get("components", {}).get("schemas", {}) + generator = ExampleGenerator(components_schemas) + paths = spec_dict.get("paths", {}) + for _path, path_item in paths.items(): + for method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]: + if method in path_item: + op = path_item[method] + if "requestBody" in op: + rb = op["requestBody"] + if "content" in rb: + for _ct, content in rb["content"].items(): + if "schema" in content: + content["example"] = generator.generate(content["schema"]) + info = spec_dict["info"] + tags = spec_dict.get("tags", []) + endpoints_by_tag: dict[str, dict[str, dict[str, Any]]] = {} + for path, path_item in paths.items(): + for method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]: + if method in path_item: + op = path_item[method] + op_tags = op.get("tags", ["Other"]) + for tag in op_tags: + if tag not in endpoints_by_tag: + endpoints_by_tag[tag] = {} + if path not in endpoints_by_tag[tag]: + endpoints_by_tag[tag][path] = {} + endpoints_by_tag[tag][path][method] = op + servers = spec_dict.get("servers", []) + components = spec_dict.get("components", {}) + output = template.render( + spec=spec_dict, + info=info, + paths=paths, + servers=servers, + tags=tags, + endpoints_by_tag=endpoints_by_tag, + components=components, + security=spec_dict.get("security", []), + external_docs=spec_dict.get("externalDocs"), + ) + Path(output_path).write_text(output) + + +def generate_markdown(spec_path: str, output_path: str, template_path: str | None = None) -> None: + spec = parse_openapi_spec(spec_path) + if template_path: + template_dir = Path(template_path).parent + else: + template_dir = Path(__file__).parent.parent / "templates" + loader = Jinja2Loader(template_dir) + env = jinja2.Environment(loader=loader) + env.filters["tojson"] = lambda x: json.dumps(x, indent=2) + template = env.get_template( + Path(template_path).name if template_path else "markdown_template.md" + ) + spec_dict = spec.model_dump() + components_schemas = spec_dict.get("components", {}).get("schemas", {}) + generator = ExampleGenerator(components_schemas) + paths = spec_dict.get("paths", {}) + for _path, path_item in paths.items(): + for method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]: + if method in path_item: + op = path_item[method] + if "requestBody" in op: + rb = op["requestBody"] + if "content" in rb: + for _ct, content in rb["content"].items(): + if "schema" in content: + content["example"] = generator.generate(content["schema"]) + info = spec_dict["info"] + tags = spec_dict.get("tags", []) + endpoints_by_tag: dict[str, dict[str, dict[str, Any]]] = {} + for path, path_item in paths.items(): + for method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]: + if method in path_item: + op = path_item[method] + op_tags = op.get("tags", ["Other"]) + for tag in op_tags: + if tag not in endpoints_by_tag: + endpoints_by_tag[tag] = {} + if path not in endpoints_by_tag[tag]: + endpoints_by_tag[tag][path] = {} + endpoints_by_tag[tag][path][method] = op + servers = spec_dict.get("servers", []) + components = spec_dict.get("components", {}) + output = template.render( + spec=spec_dict, + info=info, + paths=paths, + servers=servers, + tags=tags, + endpoints_by_tag=endpoints_by_tag, + components=components, + security=spec_dict.get("security", []), + external_docs=spec_dict.get("externalDocs"), + ) + Path(output_path).write_text(output) + + +def generate_json(spec_path: str, output_path: str, template_path: str | None = None) -> None: + spec = parse_openapi_spec(spec_path) + spec_dict = spec.model_dump() + components_schemas = spec_dict.get("components", {}).get("schemas", {}) + generator = ExampleGenerator(components_schemas) + paths = spec_dict.get("paths", {}) + for _path, path_item in paths.items(): + for method in ["get", "put", "post", "delete", "options", "head", "patch", "trace"]: + if method in path_item: + op = path_item[method] + if "requestBody" in op: + rb = op["requestBody"] + if "content" in rb: + for _ct, content in rb["content"].items(): + if "schema" in content: + content["example"] = generator.generate(content["schema"]) + for _status_code, _response in spec_dict.get("paths", {}).items(): + pass + output = json.dumps(spec_dict, indent=2) + Path(output_path).write_text(output) + + +class LocalDocsHandler(http.server.SimpleHTTPRequestHandler): + def __init__(self, *args, directory: str | None = None, **kwargs): + self.docs_dir = directory + super().__init__(*args, directory=directory, **kwargs) + + def do_GET(self): + if self.path == "/": + self.path = "/index.html" + return super().do_GET() + + +class _LocalDocsHandlerWithDir(LocalDocsHandler): + def __init__(self, *args, directory: str, **kwargs): + self.docs_dir = directory + super().__init__(*args, directory=directory, **kwargs) + + +def serve_docs(spec_path: str, host: str = "127.0.0.1", port: int = 8080) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + generate_html(spec_path, str(Path(tmpdir) / "index.html")) + shutil.copy(Path(__file__).parent.parent / "templates" / "html_template.html", tmpdir) + try: + Path.cwd().chdir(tmpdir) + with socketserver.TCPServer( + (host, port), + lambda *args, **kwargs: _LocalDocsHandlerWithDir(*args, directory=tmpdir, **kwargs) + ) as httpd: + click.echo(f"Serving API documentation at http://{host}:{port}") + httpd.serve_forever() + except KeyboardInterrupt: + pass