~abentley/bzrtools/bzrtools.dev

from dotgraph import Node, dot_output, invoke_dot, NoDot, mail_map

from bzrlib.branch import Branch

from bzrlib.errors import BzrCommandError

import bzrlib.errors

import re

import os.path

mail_map.update({'aaron.bentley@utoronto.ca'     : 'Aaron Bentley',

                 'abentley@panoramicfeedback.com': 'Aaron Bentley',

                 'john@arbash-meinel.com'        : 'John A. Meinel',

                 'mbp@sourcefrog.net'            : 'Martin Pool'

                })

def add_relations(rev_id):

    if rev_id in ancestors:

        return

    print rev_id

    if rev_id not in nodes:

        nodes[rev_id] = Node("n%d" % counter, label = rev_id)

        counter += 1

    revision = branch.get_revision(rev_id)

    ancestors [rev_id] = []

    for p in (p.revision_id for p in revision.parents):

        add_relations(p)

        if p not in descendants:

            descendants[p] = []

        descendants[p].append(rev_id)

        ancestors [rev_id].append(rev_id)

def short_committer(committer):

    new_committer = re.sub('<.*>', '', committer).strip(' ')

    if len(new_committer) < 2:

        return committer

    return new_committer

def compact_descendants(descendants, ancestors):

    new_descendants={}

    skip = set()

    for me, my_descendants in descendants.iteritems():

        if me in skip:

            continue

        new_descendants[me] = []

        for descendant in my_descendants:

            new_descendant = descendant

            while new_descendant in descendants and \

                len(ancestors[new_descendant]) == 1 and \

                len(descendants[new_descendant]) == 1:

                skip.add(new_descendant)

                if new_descendant in new_descendants:

                    del new_descendants[new_descendant]

                new_descendant = descendants[new_descendant][0]

            new_descendants[me].append(new_descendant)

    return new_descendants    

def graph_ancestry(branch, collapse=True):

    nodes = {}

    q = ((i+1, n) for (i, n) in enumerate(branch.revision_history()))

    r = 1

    try:

        branch_name = os.path.basename(branch.base)

    except AttributeError:

        branch_name = "main"

    for (revno, rev_id) in q:

        nodes[rev_id] = Node("R%d" % revno, color="#ffff00", rev_id=rev_id, 

                             cluster=branch_name)

    ancestors = {} 

    descendants = {}

    counter = 0

    lines = [branch.last_patch()]

    while len(lines) > 0:

        new_lines = set()

        for rev_id in lines:

            if rev_id not in nodes:

                nodes[rev_id] = Node("n%d" % counter, label=rev_id, 

                                     rev_id=rev_id)

                counter+=1

            try:

                revision = branch.get_revision(rev_id)

            except bzrlib.errors.NoSuchRevision:

                nodes[rev_id].node_style.append('dotted')

                continue

            if nodes[rev_id].committer is None:

                nodes[rev_id].committer = short_committer(revision.committer)

            parent_ids = [r.revision_id for r in revision.parents]

            ancestors [rev_id] = parent_ids

            for parent in parent_ids:

                if parent not in ancestors:

                    new_lines.add(parent)

                    descendants[parent] = []

                descendants[parent].append(rev_id)

        lines = new_lines

    node_relations = []

    if collapse:

        visible_descendants = compact_descendants(descendants, ancestors)

    else:

        visible_descendants = descendants

    for key, values in visible_descendants.iteritems():

        for value in values:

            node_relations.append((nodes[key], nodes[value]))

    return node_relations

def write_ancestry_file(branch, filename, collapse=True):

    b = Branch(branch)

    relations = graph_ancestry(b, collapse)

    ext = filename.split('.')[-1]

    if ext in ('svg', 'svgz', 'gif', 'jpg', 'ps', 'fig', 'mif', 'png'):

        try:

            invoke_dot(dot_output(relations), filename, ext)

        except NoDot, e:

            raise BzrCommandError("Can't find 'dot'.  Please ensure Graphviz"\

                " is installed correctly.")

    elif ext=='dot':

        file(filename, 'wb').write("".join(list(dot_output(relations))))

    else:

        print "Unknown file extension: %s" % ext