Filter OpenAPI spec by one or more paths; keep related info like components, servers, etc

filter-openapi-spec.py

import sys
import json
import argparse


def find_refs(data, key="$ref"):
    """Recursively search for values of a specified key in a nested dictionary or list.
    """
    if isinstance(data, dict):
        for k, v in data.items():
            if k == key:
                yield v
            else:
                yield from find_refs(v, key)
    elif isinstance(data, list):
        for item in data:
            yield from find_refs(item, key)


def build_ref_map(openapi_spec, filtered_paths):
    """For each schema reference found in our 'filtered' spec,
    find its definition in the 'components' top level node
    and return a map/dict.
    """

    component_map = {}

    for r in list(find_refs(filtered_paths)):
        splode = r.split('/')
        cat = splode[2]
        key = splode[3]
        item = openapi_spec.get('components',{}).get(cat,{}).get(key,{})

        if item == {}:
            print("Cannot find reference in openapi spec file.. can't proceed.")
            exit(1)

        if component_map.get(cat,None) == None:
            component_map[cat] = {}

        component_map[cat][key] = item

    return component_map

    

def filter_openapi_spec(input_file, output_file, prefixes):
    with open(input_file, 'r') as file:
        spec = json.load(file)

    # Filter paths
    filtered_paths = {path: info for path, info in spec['paths'].items() if any(
        path.startswith(prefix) for prefix in prefixes)}

    # Find initial list of $refs from our paths of interest
    components = build_ref_map(spec, filtered_paths)
    
    # Construct new spec with filtered paths
    new_spec = {
        "openapi": spec["openapi"],
        "info": spec["info"],
        "servers": spec["servers"],
        "paths": filtered_paths,
        # Include other necessary keys like 'components' if needed
        "components": components
    }

    # Recursively find schema references in our new spec, add them to the new spec, then check whether any
    # new refs have been found.. break when we flatline.
    while True:
        ref_count = set(list(find_refs(new_spec)))

        new_comps = build_ref_map(spec, new_spec)
        components = { **components, **new_comps }
        new_spec['components'] = components

        new_ref_count = set(list(find_refs(new_spec)))

        if new_ref_count == ref_count:
            break

    # Manually inject the securitySchemes.. todo: automate?
    new_spec['components']['securitySchemes'] = spec.get('components').get('securitySchemes')

    # Write the new spec to file
    with open(output_file, 'w') as file:
        json.dump(new_spec, file, indent=2)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_spec_file', '-i', 
        type=str, 
        required=True,
        help='OpenAPI spec to process')
    parser.add_argument('--output_spec_file', '-o', 
        type=str, 
        required=True,
        help='OpenAPI spec to process')
    parser.add_argument('--path_prefix', '-p', 
        type=str, 
        action='append', 
        required=True,
        help='path to extract')

    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()

    if args.path_prefix is None:
        print("No paths provided - not cool.")
        sys.exit(1)
    else:
        filter_openapi_spec(
            args.input_spec_file, 
            args.output_spec_file, 
            args.path_prefix)



if __name__ == "__main__":
    main()