soleprint/soleprint/atlas2/docgen/extractors/openapi.py

"""
An OpenAPI document -> IR: the endpoints, and the shapes they carry.

Adapts `modelgen/loader/extract/openapi.py`, which already parses OpenAPI 3.x
and Swagger 2.0 and resolves `$ref`. This turns its output into IR nodes; it
does not re-parse anything.

    python3 -m docgen extract openapi --spec petstore.yaml -o ir.json

## Why this one matters more than it looks

A hand-written API notebook is the thing nobody can keep current: the spec moves
and the document does not, and there is no way to tell by looking. Extracting
the endpoints means the document is **generated from the same file the server is
built from**, so "is this current" becomes a question about a build rather than
about somebody's diligence.

It emits schemas as `table`/`column`, the same vocabulary the database extractor
uses. That is deliberate: an API's data model and a database's are the same kind
of thing, so the ER emitter draws either without knowing which it got. Endpoints
are a separate `kind`, so a view can ask for one or the other.

    only_kinds(ir, {"table", "column"})   -> the data model, as an ER diagram
    only_kinds(ir, {"endpoint"})          -> the surface, as a notebook
"""

from pathlib import Path

from ..ir import Graph, Meta


def _modelgen():
    """modelgen's OpenAPI reader, from wherever the reference repo is.

    Imported lazily and by path rather than as a hard dependency. This is the
    **only** seam between docgen and the wider repo — see `reference.py`, which
    resolves it from `$DOCGEN_REFERENCE` or by walking up. Every other extractor
    and every emitter works with nothing above `docgen/`.

    Not reimplemented here on purpose: modelgen already parses the spec and
    resolves `$ref`, and a second OpenAPI reader in the same repo is two things
    to keep correct.
    """
    from .. import reference

    if reference.on_path() is None:
        raise reference.missing(
            "modelgen",
            "OpenAPI is read through station/tools/modelgen/loader/extract/"
            "openapi.py, which parses the spec and resolves $ref",
        )
    try:
        from modelgen.loader.extract.openapi import OpenAPIExtractor
    except ImportError as e:
        raise reference.missing(
            "modelgen.loader.extract.openapi",
            f"the reference repo was found but the module did not import ({e})",
        ) from None
    return OpenAPIExtractor


def _type_name(hint) -> str:
    if hint is None:
        return "Any"
    if isinstance(hint, str):
        return hint
    return getattr(hint, "__name__", str(hint))


def _refs(path: Path) -> dict[str, dict[str, str]]:
    """{schema: {field: referenced schema}} — the relationships, recovered.

    modelgen resolves an inter-schema `$ref` to the literal string `dict`, so by
    the time its fields reach us the *target* is gone. Without this, an API's
    data model draws as disconnected cards: three tables, no foreign keys, and
    nothing saying the relationships were lost. Which is exactly the failure
    this tool is otherwise built to prevent.

    So one key is read directly, and one only: `$ref` under a schema's
    `properties`. That is not parsing OpenAPI — no paths, no bodies, no
    responses, no `$ref` resolution, no composition keywords. modelgen still
    does all of the reading that matters, and this recovers the single fact its
    type mapping cannot carry.

    Returns `{}` on anything unexpected. A missing relationship is a worse
    diagram; a raised exception here would be no diagram at all.
    """
    try:
        import yaml                          # available: modelgen just used it
    except ImportError:
        return {}
    try:
        doc = yaml.safe_load(path.read_text()) or {}
        schemas = ((doc.get("components") or {}).get("schemas")) or {}
    except Exception:                        # noqa: BLE001 - see the docstring
        return {}

    out: dict[str, dict[str, str]] = {}
    for name, schema in schemas.items():
        if not isinstance(schema, dict):
            continue
        for field, spec in (schema.get("properties") or {}).items():
            if not isinstance(spec, dict):
                continue
            # A direct reference, or an array of them — `lines: [OrderLine]` is
            # the same relationship as `order: Order`, pointing the other way.
            ref = spec.get("$ref")
            if not ref and isinstance(spec.get("items"), dict):
                ref = spec["items"].get("$ref")
            if isinstance(ref, str) and ref.startswith("#/components/schemas/"):
                out.setdefault(name, {})[field] = ref.rsplit("/", 1)[-1]
    return out


def extract(spec_path, source: str = "openapi", identity: str | None = None) -> Graph:
    """An OpenAPI file -> IR."""
    from ..book.larder import Larder

    OpenAPIExtractor = _modelgen()
    path = Path(spec_path)
    extractor = OpenAPIExtractor(path)
    models, enums = extractor.extract()
    endpoints = extractor.endpoints()

    g = Graph(Meta(source=source, root=path.name))
    known = {m.name for m in models}
    refs = _refs(path)

    for model in models:
        attrs = {}
        if getattr(model, "docstring", None):
            attrs["doc"] = model.docstring.strip().split("\n")[0]
        g.node(model.name, "table", model.name, attrs=attrs)
        for field in model.fields:
            a = {"type": _type_name(field.type_hint)}
            if getattr(field, "optional", False):
                a["nullable"] = True
            if field.name in ("id", "uuid"):
                a["pk"] = True
            # The recovered $ref target wins over the mapped type name: modelgen
            # says `dict` where the spec said which schema.
            target = refs.get(model.name, {}).get(field.name) or _type_name(field.type_hint)
            if target in known and target != model.name:
                a["references"] = target
                g.edge(model.name, target, "foreign_key", attrs={"label": field.name})
            g.node(
                f"{model.name}.{field.name}", "column", field.name,
                parent=model.name, attrs=a,
            )

    for e in endpoints:
        eid = f"{e.method} {e.path}"
        attrs = {
            "method": e.method,
            "path": e.path,
            "status": getattr(e, "status", None) or 200,
        }
        for key in ("summary", "operation_id", "envelope_key"):
            value = getattr(e, key, None)
            if value:
                attrs[key] = value
        if getattr(e, "response_is_list", False):
            attrs["returns_list"] = True
        if getattr(e, "path_params", None):
            attrs["path_params"] = list(e.path_params)
        if getattr(e, "request_model", None):
            attrs["request_model"] = e.request_model
        if getattr(e, "response_model", None):
            attrs["response_model"] = e.response_model

        g.node(eid, "endpoint", eid, attrs=attrs)
        # `accepts` and `returns` rather than one `uses`: which direction a
        # shape travels is the thing a reader wants to know.
        if getattr(e, "request_model", None) in known:
            g.edge(eid, e.request_model, "accepts")
        if getattr(e, "response_model", None) in known:
            g.edge(eid, e.response_model, "returns")

    # A spec is the larder here, and `path` is the unit because that is what a
    # spec is an inventory of. Schemas and enums are counted separately: a spec
    # with 40 schemas and 3 endpoints is a data model, and the measure should
    # make that visible before the diagram does.
    larder = Larder(kind="openapi", identity=identity or path.name, unit="path",
                    seen=len({e.path for e in endpoints}))
    larder.extra["operations"] = len(endpoints)
    larder.extra["schemas"] = len(models)
    if enums:
        larder.extra["enums"] = len(enums)
    g.meta.larder = larder.to_dict()
    return g