scripts/cros_generate_deps_graphs.py

# Copyright 2010 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Generates pretty dependency graphs for Chrome OS packages."""

import json
import os
import sys

from chromite.lib import commandline
from chromite.lib import dot_helper


NORMAL_COLOR = "black"
TARGET_COLOR = "red"
SEED_COLOR = "green"
CHILD_COLOR = "grey"


def GetReverseDependencyClosure(full_name, deps_map):
    """Gets the closure of the reverse dependencies of a node.

    Walks the tree along all the reverse dependency paths to find all the nodes
    that transitively depend on the input node.
    """
    s = set()

    def GetClosure(name) -> None:
        s.add(name)
        node = deps_map[name]
        for dep in node["rev_deps"]:
            if dep in s:
                continue
            GetClosure(dep)

    GetClosure(full_name)
    return s


def GetOutputBaseName(node, options):
    """Gets the basename of the output file for a node."""
    return "%s_%s-%s.%s" % (
        node["category"],
        node["name"],
        node["version"],
        options.format,
    )


def AddNodeToSubgraph(subgraph, node, options, color) -> None:
    """Gets the dot definition for a node."""
    name = node["full_name"]
    href = None
    if options.link:
        filename = GetOutputBaseName(node, options)
        href = "%s%s" % (options.base_url, filename)
    subgraph.AddNode(name, name, color, href)


def GenerateDotGraph(package, deps_map, options):
    """Generates the dot source for the dependency graph leading to a node.

    The output is a list of lines.
    """
    deps = GetReverseDependencyClosure(package, deps_map)
    node = deps_map[package]

    # Keep track of all the emitted nodes so that we don't issue multiple
    # definitions
    emitted = set()

    graph = dot_helper.Graph(package)

    # Add all the children if we want them, all of them in their own subgraph,
    # as a sink. Keep the arcs outside of the subgraph though (it generates
    # better layout).
    children_subgraph = None
    if options.children and node["deps"]:
        children_subgraph = graph.AddNewSubgraph("sink")
        for child in node["deps"]:
            child_node = deps_map[child]
            AddNodeToSubgraph(
                children_subgraph, child_node, options, CHILD_COLOR
            )
            emitted.add(child)
            graph.AddArc(package, child)

    # Add the package in its own subgraph. If we didn't have children, make it
    # a sink
    if children_subgraph:
        rank = "same"
    else:
        rank = "sink"
    package_subgraph = graph.AddNewSubgraph(rank)
    AddNodeToSubgraph(package_subgraph, node, options, TARGET_COLOR)
    emitted.add(package)

    # Add all the other nodes, as well as all the arcs.
    for dep in deps:
        dep_node = deps_map[dep]
        if not dep in emitted:
            color = NORMAL_COLOR
            if dep_node["action"] == "seed":
                color = SEED_COLOR
            AddNodeToSubgraph(graph, dep_node, options, color)
        for j in dep_node["rev_deps"]:
            graph.AddArc(j, dep)

    return graph.Gen()


def GenerateImages(data, options) -> None:
    """Generate the output images for all the nodes in the input."""
    deps_map = json.loads(data)

    for package in deps_map:
        lines = GenerateDotGraph(package, deps_map, options)

        filename = os.path.join(
            options.output_dir, GetOutputBaseName(deps_map[package], options)
        )

        save_dot_filename = None
        if options.save_dot:
            save_dot_filename = filename + ".dot"

        dot_helper.GenerateImage(
            lines, filename, options.format, save_dot_filename
        )


def GetParser():
    """Return a command line parser."""
    parser = commandline.ArgumentParser(description=__doc__)
    parser.add_argument(
        "-f",
        "--format",
        default="svg",
        help="Dot output format (png, svg, etc.).",
    )
    parser.add_argument(
        "-o", "--output-dir", default=".", help="Output directory."
    )
    parser.add_argument(
        "-c", "--children", action="store_true", help="Also add children."
    )
    parser.add_argument(
        "-l", "--link", action="store_true", help="Embed links."
    )
    parser.add_argument(
        "-b", "--base-url", default="", help="Base url for links."
    )
    parser.add_argument(
        "-s", "--save-dot", action="store_true", help="Save dot files."
    )
    parser.add_argument("inputs", nargs="*", help="Chromium OS package lists")
    return parser


def main(argv) -> None:
    parser = GetParser()
    options = parser.parse_args(argv)
    options.Freeze()

    try:
        os.makedirs(options.output_dir)
    except OSError:
        # The directory already exists.
        pass

    if not options.inputs:
        GenerateImages(sys.stdin.read(), options)
    else:
        for i in options.inputs:
            with open(i, encoding="utf-8") as handle:
                GenerateImages(handle.read(), options)