path: root/Library/Contributions/examples/brew-graph

#!/usr/bin/env python
"""
$ brew install graphviz
$ brew graph | dot -Tsvg -ohomebrew.svg
$ open homebrew.svg
"""
from __future__ import with_statement

from contextlib import contextmanager
import re
from subprocess import Popen, PIPE
import sys


def run(command, print_command=False):
    "Run a command, returning the exit code and output."
    if print_command: print command
    p = Popen(command, stdout=PIPE)
    output, errput = p.communicate()
    return p.returncode, output


def _quote_id(id):
    return '"' + id.replace('"', '\"') + '"'


def format_attribs(attrib):
    if len(attrib) == 0:
        return ''

    values = ['%s="%s"' % (k, attrib[k]) for k in attrib]
    return '[' + ','.join(values) + ']'


class Output(object):
    def __init__(self, fd=sys.stdout, tabstyle="  "):
        self.fd = fd
        self.tabstyle = tabstyle
        self.tablevel = 0

    def close(self):
        self.fd = None

    def out(self, s):
        self.tabout()
        self.fd.write(s)

    def outln(self, s=None):
        if s is not None:
            self.tabout()
            self.fd.write(s)
        self.fd.write('\n')

    @contextmanager
    def indented(self):
        self.indent()
        yield self
        self.dedent()

    def indent(self):
        self.tablevel += 1

    def dedent(self):
        if self.tablevel == 0:
            raise Exception('No existing indent level.')
        self.tablevel -= 1

    def tabout(self):
        if self.tablevel:
            self.fd.write(self.tabstyle * self.tablevel)


class NodeContainer(object):
    def __init__(self):
        self.nodes = list()
        self.node_defaults = dict()
        # Stack of node attribs
        self._node_styles = list()

    def _node_style(self):
        if (len(self._node_styles) > 0):
            return self._node_styles[-1]
        else:
            return dict()

    def _push_node_style(self, attrib):
        self._node_styles.append(attrib)

    def _pop_node_style(self):
        return self._node_styles.pop()

    @contextmanager
    def node_styles(self, attrib):
        self._push_node_style(attrib)
        yield
        self._pop_node_style()

    def node(self, nodeid, label, attrib=None):
        _attrib = dict(self._node_style())
        if attrib is not None:
            _attrib.update(attrib)

        n = Node(nodeid, label, _attrib)
        self.nodes.append(n)
        return n

    def nodes_to_dot(self, out):
        if len(self.node_defaults) > 0:
            out.outln("node " + format_attribs(self.node_defaults) + ";")

        if len(self.nodes) == 0:
            return

        id_width = max([len(_quote_id(n.id)) for n in self.nodes])
        for node in self.nodes:
            node.to_dot(out, id_width)


class Node(object):
    def __init__(self, nodeid, label, attrib=None):
        self.id = nodeid
        self.label = label
        self.attrib = attrib if attrib is not None else dict()

    def as_dot(self, id_width=1):
        _attribs = dict(self.attrib)
        _attribs['label'] = self.label

        return '%-*s %s' % (id_width, _quote_id(self.id), format_attribs(_attribs))


    def to_dot(self, out, id_width=1):
        out.outln(self.as_dot(id_width))


class ClusterContainer(object):
    def __init__(self):
        self.clusters = list()

    def cluster(self, clusterid, label, attrib=None):
        c = Cluster(clusterid, label, self, attrib)
        self.clusters.append(c)
        return c


class Cluster(NodeContainer, ClusterContainer):
    def __init__(self, clusterid, label, parentcluster=None, attrib=None):
        NodeContainer.__init__(self)
        ClusterContainer.__init__(self)

        self.id = clusterid
        self.label = label
        self.attrib = attrib if attrib is not None else dict()
        self.parentcluster = parentcluster

    def cluster_id(self):
        return _quote_id("cluster_" + self.id)

    def to_dot(self, out):
        out.outln("subgraph %s {" % self.cluster_id())
        with out.indented():
            out.outln('label = "%s"' % self.label)
            for k in self.attrib:
                out.outln('%s = "%s"' % (k, self.attrib[k]))

            for cluster in self.clusters:
                cluster.to_dot(out)

            self.nodes_to_dot(out)
        out.outln("}")


class Edge(object):
    def __init__(self, source, target, attrib=None):
        if attrib is None:
            attrib = dict()

        self.source = source
        self.target = target
        self.attrib = attrib

    def to_dot(self, out):
        out.outln(self.as_dot())

    def as_dot(self):
        return " ".join((_quote_id(self.source), "->", _quote_id(self.target), format_attribs(self.attrib)))


class EdgeContainer(object):
    def __init__(self):
        self.edges = list()
        self.edge_defaults = dict()
        # Stack of edge attribs
        self._edge_styles = list()

    def _edge_style(self):
        if (len(self._edge_styles) > 0):
            return self._edge_styles[-1]
        else:
            return dict()

    def _push_edge_style(self, attrib):
        self._edge_styles.append(attrib)

    def _pop_edge_style(self):
        return self._edge_styles.pop()

    @contextmanager
    def edge_styles(self, attrib):
        self._push_edge_style(attrib)
        yield
        self._pop_edge_style()

    def link(self, source, target, attrib=None):
        _attrib = dict(self._edge_style())
        if attrib is not None:
            _attrib.update(attrib)

        e = Edge(source, target, _attrib)
        self.edges.append(e)
        return e

    def edges_to_dot(self, out):
        if len(self.edge_defaults) > 0:
            out.outln("edge " + format_attribs(self.edge_defaults) + ";")

        if len(self.edges) == 0:
            return

        for edge in self.edges:
            edge.to_dot(out)


class Graph(NodeContainer, EdgeContainer, ClusterContainer):
    """
    Contains the nodes, edges, and subgraph definitions for a graph to be
    turned into a Graphviz DOT file.
    """

    def __init__(self, label=None, attrib=None):
        NodeContainer.__init__(self)
        EdgeContainer.__init__(self)
        ClusterContainer.__init__(self)

        self.label = label if label is not None else "Default Label"
        self.attrib = attrib if attrib is not None else dict()

    def dot(self, fd=sys.stdout):
        try:
            self.o = Output(fd)
            self._dot()
        finally:
            self.o.close()

    def _dot(self):
        self.o.outln("digraph G {")

        with self.o.indented():
            self.o.outln('label = "%s"' % self.label)
            for k in self.attrib:
                self.o.outln('%s = "%s"' % (k, self.attrib[k]))

            self.nodes_to_dot(self.o)

            for cluster in self.clusters:
                self.o.outln()
                cluster.to_dot(self.o)

            self.o.outln()
            self.edges_to_dot(self.o)

        self.o.outln("}")


def main():
    code, output = run(["brew", "deps", "--all"])
    output = output.strip()
    depgraph = list()

    for f in output.split("\n"):
        stuff = f.split(":",2)
        name = stuff[0]
        deps = stuff[1].strip()
        if not deps:
            deps = list()
        else:
            deps = deps.split(" ")
        depgraph.append((name, deps))

    hb = Graph("Homebrew Dependencies", attrib={'labelloc':'b', 'rankdir':'LR', 'ranksep':'5'})

    used = set()
    for f in depgraph:
        for d in f[1]:
            used.add(f[0])
            used.add(d)

    for f in depgraph:
        if f[0] not in used:
            continue
        n = hb.node(f[0], f[0])
        for d in f[1]:
            hb.link(d, f[0])

    hb.dot()


if __name__ == "__main__":
    main()