Compute topological dependency graphs for a shared library. See https://www.georgevreilly.com/blog/2024/12/16/PatchingAndSplittingPythonWheels.html

deps_topo.py

#!/usr/bin/env python3

"""
Produce topological sorted graphs of shared object library dependencies.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys


def shared_object_deps(root_so: str) -> dict[str, list[str]]:
    graph: dict[str, list[str]] = {}
    lib_dir = os.path.dirname(os.path.abspath(root_so))
    queue = [os.path.basename(root_so)]
    while queue:
        item = queue.pop()
        output = subprocess.check_output(("ldd", os.path.join(lib_dir, item)), encoding="utf-8")
        items = [i for i in output.split("\n") if lib_dir in i]
        deps = [os.path.basename(i.rsplit(" => ")[-1].split()[0]) for i in items]
        # TODO: annotate with size information
        for d in deps:
            if d not in graph.setdefault(item, []):
                graph[item].append(d)
            queue.append(d)
    return graph


def topo_sort(graph: dict[str, list[str]]) -> list[str]:
    visited = set()
    ordered = []

    def visit(node):
        if node not in visited:
            visited.add(node)
            for target in graph.get(node, []):
                visit(target)
            ordered.append(node)

    for node in graph:
        visit(node)

    return ordered


def dot_graph(graph: dict[str, list[str]]) -> str:
    result = ["digraph G {"]
    for node, deps in graph.items():
        for d in sorted(deps):
            result.append(f'    "{node}" -> "{d}";')
    result.append("}")
    return "\n".join(result)


def mermaid_graph(graph: dict[str, list[str]]) -> str:
    result = [
        "```mermaid",
        "graph TD",
    ]
    node_ids = {node: f"n{idx}" for idx, node in enumerate(reversed(topo_sort(graph)))}
    for node, id in node_ids.items():
        result.append(f'    {id}[{node}]')
    for node, deps in graph.items():
        for d in sorted(deps):
            result.append(f'    {node_ids[node]} --> {node_ids[d]}')
    result.append("```")
    return "\n".join(result)


# From unzip -d ~/torch212 ~/torch-2.1.2+cu118_stripe.3-cp38-cp38-linux_x86_64.whl
SO_LIB = os.path.expanduser("~/torch212/torch/lib/libtorch_cuda_linalg.so")


if __name__ == "__main__":
    so_lib = sys.argv[1] if len(sys.argv) >= 2 else SO_LIB
    graph = shared_object_deps(so_lib)
    json_blob = json.dumps(graph, indent=4)

    print(f"JSON Dependency graph for {so_lib}")
    print(json_blob)

    print("\n\nDOT (Graphviz):")
    print(dot_graph(graph))

    print("\n\nMermaid graph:")
    print(mermaid_graph(graph))

    print("\n\nTopological order:")
    print(topo_sort(graph))