[2024-01-31] Script to prune GraphQL fragments based on dependency of other fragments

prune_graphql_fragments.py

import graphql.language
import networkx as nx
import tabulate


def _find_field_fragment_dependencies(
    field_node: graphql.language.ast.FieldNode,
) -> list[str]:
    if field_node.selection_set is None:
        return []

    dependencies: list[str] = []
    for selection in field_node.selection_set.selections:
        if isinstance(selection, graphql.language.ast.FieldNode):
            dependencies.extend(_find_field_fragment_dependencies(selection))
            continue

        if isinstance(selection, graphql.language.ast.FragmentSpreadNode):
            dependencies.append(selection.name.value)
            continue

        raise TypeError("Unhandled selection", selection)

    return []


def find_fragment_dependencies(
    fragment_node: graphql.language.ast.FragmentDefinitionNode,
) -> list[str]:
    dependencies: list[str] = []

    selections = fragment_node.selection_set.selections
    for selection in selections:
        if isinstance(selection, graphql.language.ast.FieldNode):
            dependencies.extend(_find_field_fragment_dependencies(selection))
            continue

        if isinstance(selection, graphql.language.ast.FragmentSpreadNode):
            dependencies.append(selection.name.value)
            continue

        raise TypeError("Unhandled selection", selection)

    return dependencies


def main():
    with open("combined.graphql", "r") as file_obj:
        content = file_obj.read()

    tree = graphql.language.parse(content)
    graph = nx.DiGraph()

    for fragment in tree.definitions:
        if not isinstance(fragment, graphql.language.ast.FragmentDefinitionNode):
            continue

        fragment_name = fragment.name.value
        graph.add_node(fragment_name)

        dependencies = find_fragment_dependencies(fragment)
        for dependency in dependencies:
            graph.add_edge(fragment_name, dependency)

    table = []
    for fragment_name, out_degree in graph.out_degree():
        if out_degree != 0:
            continue

        in_degree = graph.in_degree(fragment_name)
        table.append((fragment_name, in_degree))

    print(
        tabulate.tabulate(
            table, headers=("Fragment", "Depended on by"), tablefmt="psql"
        )
    )


if __name__ == "__main__":
    main()