apt-tree

apt-tree

#!/usr/bin/env python3
"""Display a graph of dependencies between Ubuntu packages.
Reads a list of packages from either a file or stdin.
Outputs a PDF.
"""
__author__ = "Martin Blais <[email protected]>"

import argparse
import logging
import random
import re
import sys
import tempfile

from concurrent import futures
import networkx as nx
import subprocess

def make_file(ext):
    return tempfile.mkstemp(prefix="apt-tree.", suffix=ext)[1]

def main():
    parser = argparse.ArgumentParser(description=__doc__.strip())
    parser.add_argument(
        "--filename", "-f", help="A filename with package names, one per line."
    )
    parser.add_argument(
        "--output", "-o", help="Output filename."
    )
    parser.add_argument(
        "--output-pdf", "-p", nargs="?", default=False, help="PDF output filename."
    )
    parser.add_argument("packages", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    packages = args.packages
    if args.filename:
        packages.extend(x.strip() for x in open(args.filename).readlines())
    elif not args.packages:
        packages.extend(x.strip() for x in sys.stdin.readlines())

    random.shuffle(packages)

    g = nx.DiGraph()
    with futures.ProcessPoolExecutor() as executor:
        for package, deps in zip(packages, executor.map(get_dependencies, packages)):
            g.add_node(package)
            g.add_nodes_from(deps)
            g.add_edges_from((package, dep) for dep in deps)

    outfile = args.output or make_file(".dot")
    nx.nx_agraph.write_dot(g, outfile)
    if not args.output:
        print(f"Wrote: '{outfile}'")

    if not isinstance(args.output_pdf, bool):
        filename = args.output_pdf or make_file(".pdf")
        agraph = nx.nx_agraph.to_agraph(g)
        agraph.draw(filename, prog="dot")
        if not args.output_pdf:
            print(f"Wrote: '{filename}'")


def get_dependencies(package: str):
    cp = subprocess.run(
        ["apt-cache", "depends", package], stdout=subprocess.PIPE, encoding="utf8"
    )
    if cp.returncode != 0:
        return {}
    else:
        deps = []
        for line in cp.stdout.splitlines():
            match = re.search(r".*Depends: (.*)", line)
            if match:
                deps.append(match.group(1))
        return deps


if __name__ == "__main__":
    main()