~abentley/bzrtools/bzrtools.dev

from dotgraph import Node, dot_output, invoke_dot, invoke_dot_aa, NoDot, NoRsvg

from dotgraph import mail_map, RSVG_OUTPUT_TYPES, DOT_OUTPUT_TYPES, Edge

from bzrlib.branch import Branch

from bzrlib.errors import BzrCommandError, NoCommonRoot, NoSuchRevision

from bzrlib.fetch import greedy_fetch

from bzrlib.graph import node_distances, select_farthest

from bzrlib.revision import combined_graph, MultipleRevisionSources

import bzrlib.errors

import re

import os.path

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

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

                 'abentley@lappy'                : 'Aaron Bentley',

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

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

                 'robertc@robertcollins.net'     : 'Robert Collins',

                })

committer_alias = {'abentley': 'Aaron Bentley'}

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 can_skip(rev_id, descendants, ancestors):

    if rev_id not in descendants:

        return False

    elif rev_id not in ancestors:

        return False

    elif len(ancestors[rev_id]) != 1:

        return False

    elif len(descendants[list(ancestors[rev_id])[0]]) != 1:

        return False

    elif len(descendants[rev_id]) != 1:

        return False

    else:

        return True

def compact_ancestors(descendants, ancestors, exceptions=()):

    new_ancestors={}

    skip = set()

    for me, my_parents in ancestors.iteritems():

        if me in skip:

            continue

        new_ancestors[me] = {} 

        for parent in my_parents:

            new_parent = parent 

            distance = 0

            while can_skip(new_parent, descendants, ancestors):

                if new_parent in exceptions:

                    break

                skip.add(new_parent)

                if new_parent in new_ancestors:

                    del new_ancestors[new_parent]

                new_parent = list(ancestors[new_parent])[0]

                distance += 1

            new_ancestors[me][new_parent] = distance

    return new_ancestors    

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 can_skip(new_descendant, descendants, ancestors):

                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 = []

    for node in nodes.itervalues():

        node.label = node.get_label()

    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 get_committer(rev_id, source):

    try:

        committer = short_committer(source.get_revision(rev_id).committer)

    except NoSuchRevision:

        try:

            committer = '-'.join(rev_id.split('-')[:-2]).strip(' ')

            if committer == '':

                return None

        except ValueError:

            return None

    if '@' in committer:

        try:

            committer = mail_map[committer]

        except KeyError:

            pass

    try:

        committer = committer_alias[committer]

    except KeyError:

        pass

    return committer

def graph_merge_pick(branch, other_branch, collapse=False):

    greedy_fetch(branch, other_branch)

    revision_a = branch.last_patch()

    revision_b = other_branch.last_patch()

    try:

        root, ancestors, descendants, common = \

            combined_graph(revision_a, revision_b, branch)

    except bzrlib.errors.NoCommonRoot:

        raise bzrlib.errors.NoCommonAncestor(revision_a, revision_b)

    distances = node_distances(descendants, ancestors, root)

    base = select_farthest(distances, common)

    n_history = branch.revision_history()

    m_history = other_branch.revision_history() 

    dot_nodes = {}

    def dot_node(node, num):

        try:

            n_rev = n_history.index(node) + 1

        except ValueError:

            n_rev = None

        try:

            m_rev = m_history.index(node) + 1

        except ValueError:

            m_rev = None

        if (n_rev, m_rev) == (None, None):

            name = node[-5:]

            cluster = None

        elif n_rev == m_rev:

            name = "rR%d" % n_rev

        else:

            namelist = []

            for prefix, revno in (('r', n_rev), ('R', m_rev)):

                if revno is not None:

                    namelist.append("%s%d" % (prefix, revno))

            name = ' '.join(namelist)

        if None not in (n_rev, m_rev):

            cluster = "common_history"

            color = "#ff9900"

        elif (None, None) == (n_rev, m_rev):

            cluster = None

            if node in common:

                color = "#6699ff"

            else:

                color = None

        elif n_rev is not None:

            cluster = "my_history"

            color = "#ffff00"

        else:

            assert m_rev is not None

            cluster = "other_history"

            color = "#ff0000"

        if node == base:

            color = "#33ff99"

        label = [name]

        committer = get_committer(node, branch)

        if committer is not None:

            label.append(committer)

        if node in distances:

            label.append('d%d' % distances[node])

        d_node = Node("n%d" % num, color=color, label="\\n".join(label), 

                    rev_id=node, cluster=cluster)

        if node not in ancestors:

            d_node.node_style.append('dotted')

        return d_node

    node_relations = []

    num = 0

    if collapse:

        visible_ancestors = compact_ancestors(descendants, ancestors, (base,))

    else:

        visible_ancestors = ancestors

    for node, parents in visible_ancestors.iteritems():

        if node not in dot_nodes:

            dot_nodes[node] = dot_node(node, num)

            num += 1

        if visible_ancestors is ancestors:

            parent_iter = ((f, 0) for f in parents)

        else:

            parent_iter = (f for f in parents.iteritems())

        for parent, skipped in parent_iter:

            if parent not in dot_nodes:

                dot_nodes[parent] = dot_node(parent, num)

                num += 1

            edge = Edge(dot_nodes[parent], dot_nodes[node])

            if skipped != 0:

                edge.label = "%d" % skipped

            node_relations.append(edge)

    return node_relations

def write_ancestry_file(branch, filename, collapse=True, antialias=True,

                        merge_branch=None):

    b = Branch.open_containing(branch)

    if merge_branch is None:

        relations = graph_ancestry(b, collapse)

    else:

        m = Branch.open_containing(merge_branch)

        relations = graph_merge_pick(b, m, collapse)

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

    if antialias and ext in RSVG_OUTPUT_TYPES:

        try:

            invoke_dot_aa(dot_output(relations), filename, ext)

        except NoDot, e:

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

                " is installed correctly.")

        except NoRsvg, e:

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

                "librsvg-bin is installed correctly, or use --noantialias.")

    elif ext in DOT_OUTPUT_TYPES:

        try:

            invoke_dot(dot_output(relations), filename, ext)

        except NoDot, e:

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

                " is installed correctly, or use --noantialias")

    elif ext=='dot':

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

    else:

        print "Unknown file extension: %s" % ext