combine a local openapi.yaml with a remote definitions file, to eliminate the remote refs

eliminate_remote_refs.py

from typing import Any

import requests
import ruamel.yaml



def remove_unused(d: Any) -> None:
    """
    trims all schemas in components/schemas that are not referenced in the spec
    """
    used_schemas = set()
    used_params = set()

    def _remove_unused(d: Any) -> None:
        if isinstance(d, list):
            for i in range(len(d)):
                _remove_unused(d[i])

        if isinstance(d, dict):
            for k in d:
                if k == "$ref":
                    parts = d[k].split("/")[-2:]
                    if parts[0] == "schemas":
                        used_schemas.add(parts[-1])
                    elif parts[0] == "parameters":
                        used_params.add(parts[-1])
                else:
                    _remove_unused(d[k])

    _remove_unused(d)
    d["components"]["schemas"] = {k: v for k, v in d["components"]["schemas"].items() if k in used_schemas}
    d["components"]["parameters"] = {k: v for k, v in d["components"]["parameters"].items() if k in used_params}


def remote_refs_to_local(
    remote_url: str,
    input_fpath: str,
    output_fpath: str,
    max_line_width: int = 1000,
    indent_kwargs: dict[str, Any] | None = None,
) -> None:
    """
    remove remote urls from $ref and replace them with local references

    Note: this currently only goes "one level" if will fully resolve schema A that references schema B,
    but it will not resolve schema B if it references schema C.
    """

    yaml = ruamel.yaml.YAML()
    if indent_kwargs:
        yaml.indent(**indent_kwargs)
    else:
        yaml.indent(mapping=2, sequence=4, offset=2)
    yaml.width = max_line_width

    response = requests.get(remote_url, allow_redirects=True)
    content = yaml.load(response.content.decode("utf-8"))
    remote_params = content["components"]["parameters"]
    remote_schemas = content["components"]["schemas"]

    with open(input_fpath, "r") as f:
        spec = yaml.load(f)

    # this function relies on closures over the params
    def resolve(d: Any) -> Any:
        if isinstance(d, list):
            for i in range(len(d)):
                d[i] = resolve(d[i])

        if isinstance(d, dict):
            for k in d:
                if k == "$ref" and d[k].startswith("http"):
                    parts = d[k].split("/")[-2:]

                    if parts[0] == "parameters":
                        pname = parts[-1]
                        if pname not in remote_params:
                            raise RuntimeError("ERROR: Parameter not found", parts[-1])
                        d[k] = f"#/components/parameters/{pname}"

                    elif parts[0] == "schemas":
                        sname = parts[-1]
                        if sname not in remote_schemas:
                            raise RuntimeError("ERROR: Schema not found", parts[-1])
                        d[k] = f"#/components/schemas/{sname}"
                else:
                    d[k] = resolve(d[k])

        return d

    resolve(spec)

    if "parameters" not in spec["components"]:
        spec["components"]["parameters"] = {}
    spec["components"]["parameters"].update(remote_params)

    if "schemas" not in spec["components"]:
        spec["components"]["schemas"] = {}
    spec["components"]["schemas"].update(remote_schemas)

    # some of the remote schemas reference  their own schemas. Its easier to just add all the remove schemas, then
    # remove the ones never referenced, than it is to keep track of which schemas are referenced and only add those.
    remove_unused(spec)

    with open(output_fpath, "w") as f:
        yaml.dump(spec, f)