contrib/bst-graph

#!/usr/bin/env python3
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#  Authors:
#        Chandan Singh <csingh43@bloomberg.net>
#
'''Print dependency graph of given element(s) in DOT format.

This script must be run from the same directory where you would normally
run `bst` commands.

When `--format` option is specified, the output will also be rendered in the
given format. A file with name `bst-graph.{format}` will be created in the same
directory. To use this option, you must have the `graphviz` command line tool
installed.
'''

import argparse
import subprocess
import re
import urllib.parse

from graphviz import Digraph
from ruamel.yaml import YAML

def parse_args():
    '''Handle parsing of command line arguments.

    Returns:
       A argparse.Namespace object
    '''
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        'element', nargs='*',
        help='Name of the element'
    )
    parser.add_argument(
        '--format',
        help='Redner the graph in given format (`pdf`, `png`, `svg` etc)'
    )
    parser.add_argument(
        '--view', action='store_true',
        help='Open the rendered graph with the default application'
    )
    return parser.parse_args()


def unique_node_name(s):
    '''Generate unique node name for `s`.

    Graphviz node names cannot contain colons or backslashes so we use
    url-encoding to generate the unique node name. (A cryptographic hash could
    be used instead but that would make the graphviz file less readable.)

    Args:
       s: element name

    Returns:
       A string containing the unique node name
    '''
    return urllib.parse.quote_plus(s)


def parse_graph(lines):
    '''Return nodes and edges of the parsed grpah.

    Args:
       lines: List of lines in format 'NAME|BUILD-DEPS|RUNTIME-DEPS'

    Returns:
       Tuple of format (nodes,build_deps,runtime_deps)
       Each member of build_deps and runtime_deps is also a tuple.
    '''
    parser = YAML(typ="safe")
    nodes = set()
    build_deps = set()
    runtime_deps = set()
    for line in lines:
        line = line.strip()
        if not line:
            continue
        # It is safe to split on '|' as it is not a valid character for
        # element names.
        name, build_dep, runtime_dep = line.split('|')

        build_dep = parser.load(build_dep)
        runtime_dep = parser.load(runtime_dep)

        nodes.add(name)
        [build_deps.add((name, dep)) for dep in build_dep if dep]
        [runtime_deps.add((name, dep)) for dep in runtime_dep if dep]

    return nodes, build_deps, runtime_deps


def generate_graph(nodes, build_deps, runtime_deps):
    '''Generate graph from given nodes and edges.

    Args:
       nodes: set of nodes
       build_deps: set of tuples of build depdencies
       runtime_deps: set of tuples of runtime depdencies

    Returns:
       A graphviz.Digraph object
    '''
    graph = Digraph()
    for name in nodes:
        graph.node(unique_node_name(name), label=name)
    for source, target in build_deps:
        graph.edge(unique_node_name(source), unique_node_name(target), label='build-dep')
    for source, target in runtime_deps:
        graph.edge(unique_node_name(source), unique_node_name(target), label='runtime-dep')
    return graph


def main():
    args = parse_args()
    cmd = ['bst', 'show', '--format', '%{name}|%{build-deps}|%{runtime-deps}||']
    if 'element' in args:
        cmd += args.element
    graph_lines = subprocess.check_output(cmd, universal_newlines=True)
    # NOTE: We generate nodes and edges before giving them to graphviz as
    # the library does not de-deuplicate them.
    nodes, build_deps, runtime_deps = parse_graph(re.split(r"\|\|", graph_lines))
    graph = generate_graph(nodes, build_deps, runtime_deps)

    print(graph.source)

    if args.format:
        graph.render(cleanup=True,
                     filename='bst-graph',
                     format=args.format,
                     view=args.view)


if __name__ == '__main__':
    main()