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

#!/usr/bin/env python
"""
$ brew install graphviz
$ brew graph | dot -Tsvg -ohomebrew.html
$ open homebrew.html
"""
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():
            # If the label is an HTML-like string (starts and end with '<' and '>', respectively),
            # don't put quotes around it (or GraphViz won't recognize it.)
            if self.label[0] == '<' and self.label[-1] == '>':
                out.outln('label = %s;' % self.label)
            else:
                out.outln('label = "%s"' % self.label)
            for k in self.attrib:
                if isinstance(self.attrib[k], dict):
                    out.outln('%s %s;' % (k, format_attribs(self.attrib[k])))
                else:
                    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:
                # If the value of the attrib is a dictionary, write it out in special array form
                if isinstance(self.attrib[k], dict):
                    self.o.outln('%s %s;' % (k, format_attribs(self.attrib[k])))
                else:
                    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():
    cmd = ["brew", "deps"]
    if sys.argv[1:]:
        if '--all' in sys.argv[1:]:
            show = 'all'
            cmd.extend(['--all'])
        else:
            show = 'one'
            hideOrphaned = False
            cmd.extend(sys.argv[1:])
    else:
        show = 'installed'
        cmd.extend(['--installed'])

    code, output = run(cmd)
    output = output.strip()
    depgraph = list()

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

    # We need newrank = True to make sure clusters respect rank = "source". Otherwise, we may get
    # random nodes next to "Safe to Remove" cluster, despite them not being a part of that cluster.
    hb = Graph("Homebrew Dependencies", attrib={
        'labelloc':'t', 'rankdir': 'TB' , 'ranksep':'3', 'newrank': True,
        'graph': {'fontname': 'Futura-Medium', 'fontsize': 48},
        'node': {'fontname': 'HelveticaNeue', 'fontsize': 14}
        })
    # Independent formulas (those that are not dependended on by any other formula) get placed in
    # their own subgraph so we can align them together on the left.
    if show == 'installed':
        # We use a HTML-like label to give the label a little bit of padding at the top
        sub = hb.cluster("independent", "<<font point-size=\"15\"><br/></font>Safe to Remove>",
                attrib={'rank': 'source', 'style': 'filled', 'fillcolor': '#F0F0F0', 'color': 'invis',
                    'margin': '25,1', 'graph': {'fontname': 'Helvetica-LightOblique', 'fontsize': 24}})
    else:
        sub = hb

    seen = set()
    def addNode(graph, name):
        if name not in seen:
            graph.node(name, name, attrib={'shape': 'box'})
            seen.add(name)
            return True
        return False

    independent = set()
    for f in depgraph:
        # Filter out orphan formulas when showing all, to cut down on noise
        if show == 'all' and len(f[1]) == 0:
            continue

        independent.add(f[0])
        for d in f[1]:
            independent.discard(d)
            hb.link(f[0], d)
            # Children we can add right away because we don't care where they go
            addNode(hb, d)

    # For all installed formulas, place them in the 'indep' subgraph iff they
    # are not depended on by other formulas, i.e. are root nodes.
    for d in independent:
        addNode(sub, d)

    hb.dot()


if __name__ == "__main__":
    main()