Plot Arch packages to matplotlib or draw.io.

plot_packages.py

#!/usr/bin/env python3

import argparse
from pathlib import Path
import subprocess
import pprint
import networkx as nx
import matplotlib.pyplot as plt
from lxml import etree


def is_meta_package(package_name: str) -> bool:
    output = subprocess.check_output(["pacman", "-Ql", package_name])

    res = output.decode().split("\n")

    if len(res) == 1 and res[0] == "":
        return True

    return False


def get_package_deps(package_name: str) -> list[str]:
    output = subprocess.check_output(["pacman", "-Qi", package_name])

    packages = []
    for line in output.decode().split("\n"):
        if "Depends On" in line:
            split0 = line.split(":")
            packages = split0[1].split()

    return packages

def is_rocm_package(package: str) -> bool:
    if "hip" in package:
        return True
    if "roc" in package:
        return True

    whitelist = [
        "aotriton",
        "rccl",
        "comgr",
        "composable-kernel"
    ]

    for w in whitelist:
        if w in package:
            return True

    return False



class DepFinder:
    def __init__(self):
        self.deps_by_package = {}
        self.reject = set()

    def filter_rocm_packages(self, packages: list[str]) -> list[str]:

        res = []

        for p in packages:
            if is_rocm_package(p):
                res.append(p)
            else:
                self.reject.add(p)

        return res

    def get_rocm_deps_recusive(self, package_name: str):
        deps = get_package_deps(package_name)
        rocm_deps = self.filter_rocm_packages(deps)

        assert package_name not in self.deps_by_package

        self.deps_by_package[package_name] = rocm_deps

        # print(f"{package_name} has {len(rocm_deps)} rocm deps:", rocm_deps)

        for rocm_dep in rocm_deps:
            if rocm_dep not in self.deps_by_package:
                self.get_rocm_deps_recusive(rocm_dep)
            # res.extend(sub_res)

        # return res


def plot_nx(deps_by_package: dict):
    packages = list(deps_by_package.keys())

    dependencies = []
    for package, deps in deps_by_package.items():
        for dep in deps:
            dependencies.append((package, dep))

    G = nx.DiGraph()
    G.add_nodes_from(packages)
    G.add_edges_from(dependencies)

    # plot
    pos = nx.nx_agraph.graphviz_layout(G,
                                       prog='dot',
                                    #    prog='fdp',
                                    #    prog='neato'
                                       )

    nx.draw_networkx_nodes(G, pos, node_size=3000, node_color="#ADD8E6")
    nx.draw_networkx_edges(G, pos, edgelist=G.edges(), arrows=True, arrowstyle='-|>', arrowsize=20, edge_color="gray", width=2)
    nx.draw_networkx_labels(G, pos, font_size=10, font_weight="bold")

    plt.title("Package Dependency Graph")
    plt.axis('off')
    plt.show()



def add_edge(root, the_id: int, source: int, target: int):
    cell = etree.SubElement(root, "mxCell",
                            attrib = {
                                "id": str(the_id),
                                "parent": str(1),
                                "edge": str(1),
                                "source": str(source),
                                "target": str(target)
                            }
    )
    geometry = etree.SubElement(cell, "mxGeometry", attrib = {"relative": str(1), "as": "geometry"})


def add_node(root, the_id: int, name: str, x: int, y: int):
    cell = etree.SubElement(root, "mxCell", attrib = {
                                                "id": str(the_id),
                                                "parent": str(1),
                                                "value": name,
                                                "style": "ellipse",
                                                "vertex": str(1)
                                            }
    )
    geometry = etree.SubElement(cell, "mxGeometry", attrib = {
                                                        "x": str(x),
                                                        "y": str(y),
                                                        "width": str(220),
                                                        "height": str(100),
                                                        "as": "geometry"
                                                    }
    )

def should_ignore_package(package_name: str) -> bool:
    deny_list = [
        "rocm-core",
        "hip-runtime-amd"
    ]

    if package_name in deny_list:
        return True

    if is_meta_package(package_name):
        return True

    return False


def export_drawio(deps_by_package: dict):

    # all_packages = list(deps_by_package.keys())

    # packages = []
    # for p in all_packages:
    #     if not is_meta_package(p):
    #         packages.append(p)

    dependencies = []
    for package, deps in deps_by_package.items():
        if should_ignore_package(package):
            continue
        for dep in deps:
            if should_ignore_package(dep):
                continue
            dependencies.append((package, dep))

    package_to_graph_id = {}

    nodes_by_number_of_deps = {}

    for package_name, deps in deps_by_package.items():
        if should_ignore_package(package_name):
            continue
        number_of_deps = len(deps)
        if number_of_deps not in nodes_by_number_of_deps:
            nodes_by_number_of_deps[number_of_deps] = []
        nodes_by_number_of_deps[number_of_deps].append(package_name)

    nodes_by_number_of_deps_sorted = sorted(nodes_by_number_of_deps.items(), reverse=True)

    # pprint.pprint(nodes_by_number_of_deps_sorted)

    mx_file = etree.Element("mxfile")
    diagram = etree.SubElement(mx_file, "diagram")
    mx_graph_model = etree.SubElement(diagram, "mxGraphModel", attrib = {
        "dx": str(2066),
        "dy": str(1204),
        "grid": str(1),
        "gridSize": str(10),
        "guides": str(1),
        "tooltips": str(1),
        "connect": str(1),
        "arrows": str(1),
        "fold": str(1),
        "page": str(1),
        "pageScale": str(1),
        "pageWidth": str(850),
        "pageHeight": str(1100),
        "math": str(0),
        "shadow": str(0)
    })
    root = etree.SubElement(mx_graph_model, "root")
    mx_cell_0 = etree.SubElement(root, "mxCell", attrib = {"id": str(0)})
    mx_cell_1 = etree.SubElement(root, "mxCell", attrib = {"id": str(1), "parent": str(0)})

    current_graph_id = 1

    # nodes

    for i, (number_of_deps, package_names) in enumerate(nodes_by_number_of_deps_sorted):
        for j, package_name in enumerate(package_names):
            print(i, j, package_name)

            # current_graph_id += 1
            # package_to_graph_id["rocm-core"] = current_graph_id
            # add_node(root, current_graph_id, "rocm-core", 280, 300)

            current_graph_id += 1
            package_to_graph_id[package_name] = current_graph_id
            add_node(root, current_graph_id, package_name, j * 200 + 80, i * 200 + 80)


    # current_graph_id += 1
    # package_to_graph_id["rocm-something"] = current_graph_id
    # add_node(root, current_graph_id, "rocm-something", 280, 80)

    # edge
    # current_graph_id += 1
    # add_edge(root, current_graph_id, package_to_graph_id["rocm-core"], package_to_graph_id["rocm-something"])

    for a, b in dependencies:
        current_graph_id += 1
        add_edge(root, current_graph_id, package_to_graph_id[a], package_to_graph_id[b])

    tree = etree.ElementTree(mx_file)
    with open("python_export.xml", "wb") as f:
        tree.write(
            f,
            pretty_print=True,
            xml_declaration=True,
            encoding="utf-8"
        )

def main():
    # is_meta_package("rocm-hip-runtime")
    # is_meta_package("python-pytorch-opt-rocm")
    # return

    dep_finder = DepFinder()

    dep_finder.get_rocm_deps_recusive("python-pytorch-opt-rocm")

    # pprint.pprint(dep_finder.deps_by_package)
    # print(list(dep_finder.reject))

    # plot_nx(dep_finder.deps_by_package)
    export_drawio(dep_finder.deps_by_package)





if __name__ == "__main__":
    main()