~abentley/bzrtools/bzrtools.dev

« back to all changes in this revision

Viewing changes to graph.py

  • Committer: Aaron Bentley
  • Date: 2012-03-20 02:40:57 UTC
  • Revision ID: aaron@aaronbentley.com-20120320024057-mqltf93xxs09r0ry
Compatibility fixes for bzr.dev

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2005, 2008 Aaron Bentley
 
2
# <aaron@aaronbentley.com>
 
3
#
 
4
#    This program is free software; you can redistribute it and/or modify
 
5
#    it under the terms of the GNU General Public License as published by
 
6
#    the Free Software Foundation; either version 2 of the License, or
 
7
#    (at your option) any later version.
 
8
#
 
9
#    This program is distributed in the hope that it will be useful,
 
10
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
 
11
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
12
#    GNU General Public License for more details.
 
13
#
 
14
#    You should have received a copy of the GNU General Public License
 
15
#    along with this program; if not, write to the Free Software
 
16
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
17
 
 
18
 
 
19
import time
 
20
 
 
21
from bzrlib.branch import Branch
 
22
from bzrlib.errors import BzrCommandError, NoSuchRevision
 
23
from bzrlib.revision import NULL_REVISION
 
24
 
 
25
from bzrtools import short_committer
 
26
from dotgraph import (
 
27
    dot_output,
 
28
    DOT_OUTPUT_TYPES,
 
29
    Edge,
 
30
    invoke_dot,
 
31
    invoke_dot_aa,
 
32
    invoke_dot_html,
 
33
    Node,
 
34
    NoDot,
 
35
    NoRsvg,
 
36
    RSVG_OUTPUT_TYPES,
 
37
    )
 
38
 
 
39
 
 
40
def max_distance(node, ancestors, distances, root_descendants):
 
41
    """Calculate the max distance to an ancestor.
 
42
    Return None if not all possible ancestors have known distances"""
 
43
    best = None
 
44
    if node in distances:
 
45
        best = distances[node]
 
46
    for ancestor in ancestors[node]:
 
47
        # skip ancestors we will never traverse:
 
48
        if root_descendants is not None and ancestor not in root_descendants:
 
49
            continue
 
50
        # An ancestor which is not listed in ancestors will never be in
 
51
        # distances, so we pretend it never existed.
 
52
        if ancestor not in ancestors:
 
53
            continue
 
54
        if ancestor not in distances:
 
55
            return None
 
56
        if best is None or distances[ancestor]+1 > best:
 
57
            best = distances[ancestor] + 1
 
58
    return best
 
59
 
 
60
 
 
61
def node_distances(graph, ancestors, start, root_descendants=None):
 
62
    """Produce a list of nodes, sorted by distance from a start node.
 
63
    This is an algorithm devised by Aaron Bentley, because applying Dijkstra
 
64
    backwards seemed too complicated.
 
65
 
 
66
    For each node, we walk its descendants.  If all the descendant's ancestors
 
67
    have a max-distance-to-start, (excluding ones that can never reach start),
 
68
    we calculate their max-distance-to-start, and schedule their descendants.
 
69
 
 
70
    So when a node's last parent acquires a distance, it will acquire a
 
71
    distance on the next iteration.
 
72
 
 
73
    Once we know the max distances for all nodes, we can return a list sorted
 
74
    by distance, farthest first.
 
75
    """
 
76
    distances = {start: 0}
 
77
    lines = set([start])
 
78
    while len(lines) > 0:
 
79
        new_lines = set()
 
80
        for line in lines:
 
81
            line_descendants = graph[line]
 
82
            for descendant in line_descendants:
 
83
                distance = max_distance(descendant, ancestors, distances,
 
84
                                        root_descendants)
 
85
                if distance is None:
 
86
                    continue
 
87
                distances[descendant] = distance
 
88
                new_lines.add(descendant)
 
89
        lines = new_lines
 
90
    return distances
 
91
 
 
92
 
 
93
def nodes_by_distance(distances):
 
94
    """Return a list of nodes sorted by distance"""
 
95
    def by_distance(n):
 
96
        return distances[n],n
 
97
 
 
98
    node_list = distances.keys()
 
99
    node_list.sort(key=by_distance, reverse=True)
 
100
    return node_list
 
101
 
 
102
 
 
103
def select_farthest(distances, common):
 
104
    """Return the farthest common node, or None if no node qualifies."""
 
105
    node_list = nodes_by_distance(distances)
 
106
    for node in node_list:
 
107
        if node in common:
 
108
            return node
 
109
 
 
110
 
 
111
mail_map = {'aaron.bentley@utoronto.ca'     : 'Aaron Bentley',
 
112
            'abentley@panoramicfeedback.com': 'Aaron Bentley',
 
113
            'abentley@lappy'                : 'Aaron Bentley',
 
114
            'john@arbash-meinel.com'        : 'John Arbash Meinel',
 
115
            'mbp@sourcefrog.net'            : 'Martin Pool',
 
116
            'robertc@robertcollins.net'     : 'Robert Collins',
 
117
            }
 
118
 
 
119
committer_alias = {'abentley': 'Aaron Bentley'}
 
120
def can_skip(rev_id, descendants, ancestors):
 
121
    if rev_id not in descendants:
 
122
        return False
 
123
    elif rev_id not in ancestors:
 
124
        return False
 
125
    elif len(ancestors[rev_id]) != 1:
 
126
        return False
 
127
    elif len(descendants[list(ancestors[rev_id])[0]]) != 1:
 
128
        return False
 
129
    elif len(descendants[rev_id]) != 1:
 
130
        return False
 
131
    else:
 
132
        return True
 
133
 
 
134
def compact_ancestors(descendants, ancestors, exceptions=()):
 
135
    new_ancestors={}
 
136
    skip = set()
 
137
    for me, my_parents in ancestors.iteritems():
 
138
        if me in skip:
 
139
            continue
 
140
        new_ancestors[me] = {}
 
141
        for parent in my_parents:
 
142
            new_parent = parent
 
143
            distance = 0
 
144
            while can_skip(new_parent, descendants, ancestors):
 
145
                if new_parent in exceptions:
 
146
                    break
 
147
                skip.add(new_parent)
 
148
                if new_parent in new_ancestors:
 
149
                    del new_ancestors[new_parent]
 
150
                new_parent = list(ancestors[new_parent])[0]
 
151
                distance += 1
 
152
            new_ancestors[me][new_parent] = distance
 
153
    return new_ancestors
 
154
 
 
155
def get_rev_info(rev_id, source):
 
156
    """Return the committer, message, nick and date of a revision."""
 
157
    committer = None
 
158
    message = None
 
159
    date = None
 
160
    nick = None
 
161
    if rev_id == 'null:':
 
162
        return None, 'Null Revision', None, None
 
163
    try:
 
164
        rev = source.get_revision(rev_id)
 
165
    except NoSuchRevision:
 
166
        try:
 
167
            committer = '-'.join(rev_id.split('-')[:-2]).strip(' ')
 
168
            if committer == '':
 
169
                return None, None, None, None
 
170
        except ValueError:
 
171
            return None, None, None, None
 
172
    else:
 
173
        committer = short_committer(rev.committer)
 
174
        if rev.message is not None:
 
175
            message = rev.message.split('\n')[0]
 
176
        gmtime = time.gmtime(rev.timestamp + (rev.timezone or 0))
 
177
        date = time.strftime('%Y/%m/%d', gmtime)
 
178
        nick = rev.properties.get('branch-nick')
 
179
    if '@' in committer:
 
180
        try:
 
181
            committer = mail_map[committer]
 
182
        except KeyError:
 
183
            pass
 
184
    try:
 
185
        committer = committer_alias[committer]
 
186
    except KeyError:
 
187
        pass
 
188
    return committer, message, nick, date
 
189
 
 
190
class Grapher(object):
 
191
 
 
192
    def __init__(self, branch, other_branch=None):
 
193
        object.__init__(self)
 
194
        self.branch = branch
 
195
        self.other_branch = other_branch
 
196
        if other_branch is not None:
 
197
            other_repo = other_branch.repository
 
198
            revision_b = self.other_branch.last_revision()
 
199
        else:
 
200
            other_repo = None
 
201
            revision_b = None
 
202
        self.graph = self.branch.repository.get_graph(other_repo)
 
203
        revision_a = self.branch.last_revision()
 
204
        self.scan_graph(revision_a, revision_b)
 
205
        self.n_history = list(self.graph.iter_lefthand_ancestry(revision_a))
 
206
        self.n_history.reverse()
 
207
        self.n_revnos = branch.get_revision_id_to_revno_map()
 
208
        self.distances = node_distances(self.descendants, self.ancestors,
 
209
                                        self.root)
 
210
        if other_branch is not None:
 
211
            self.base = select_farthest(self.distances, self.common)
 
212
            self.m_history = self.graph.iter_lefthand_ancestry(revision_b)
 
213
            self.m_history = list(self.m_history)
 
214
            self.m_history.reverse()
 
215
            self.m_revnos = other_branch.get_revision_id_to_revno_map()
 
216
            self.new_base = self.graph.find_unique_lca(revision_a,
 
217
                                                       revision_b)
 
218
            self.lcas = self.graph.find_lca(revision_a, revision_b)
 
219
        else:
 
220
            self.base = None
 
221
            self.new_base = None
 
222
            self.lcas = set()
 
223
            self.m_history = []
 
224
            self.m_revnos = {}
 
225
 
 
226
    def scan_graph(self, revision_a, revision_b):
 
227
        a_ancestors = dict(self.graph.iter_ancestry([revision_a]))
 
228
        self.ancestors = a_ancestors
 
229
        self.root = NULL_REVISION
 
230
        if revision_b is not None:
 
231
            b_ancestors = dict(self.graph.iter_ancestry([revision_b]))
 
232
            self.common = set(a_ancestors.keys())
 
233
            self.common.intersection_update(b_ancestors)
 
234
            self.ancestors.update(b_ancestors)
 
235
        else:
 
236
            self.common = []
 
237
            revision_b = None
 
238
        self.descendants = {}
 
239
        ghosts = set()
 
240
        for revision, parents in self.ancestors.iteritems():
 
241
            self.descendants.setdefault(revision, [])
 
242
            if parents is None:
 
243
                ghosts.add(revision)
 
244
                parents = [NULL_REVISION]
 
245
            for parent in parents:
 
246
                self.descendants.setdefault(parent, []).append(revision)
 
247
        for ghost in ghosts:
 
248
            self.ancestors[ghost] = [NULL_REVISION]
 
249
 
 
250
    @staticmethod
 
251
    def _get_revno_str(prefix, revno_map, revision_id):
 
252
        try:
 
253
            revno = revno_map[revision_id]
 
254
        except KeyError:
 
255
            return None
 
256
        return '%s%s' % (prefix, '.'.join(str(n) for n in revno))
 
257
 
 
258
    def dot_node(self, node, num):
 
259
        try:
 
260
            n_rev = self.n_history.index(node) + 1
 
261
        except ValueError:
 
262
            n_rev = None
 
263
        try:
 
264
            m_rev = self.m_history.index(node) + 1
 
265
        except ValueError:
 
266
            m_rev = None
 
267
        if (n_rev, m_rev) == (None, None):
 
268
            name = self._get_revno_str('r', self.n_revnos, node)
 
269
            if name is None:
 
270
                name = self._get_revno_str('R', self.m_revnos, node)
 
271
            if name is None:
 
272
                name = node[-5:]
 
273
            cluster = None
 
274
        elif n_rev == m_rev:
 
275
            name = "rR%d" % n_rev
 
276
        else:
 
277
            namelist = []
 
278
            for prefix, revno in (('r', n_rev), ('R', m_rev)):
 
279
                if revno is not None:
 
280
                    namelist.append("%s%d" % (prefix, revno))
 
281
            name = ' '.join(namelist)
 
282
        if None not in (n_rev, m_rev):
 
283
            cluster = "common_history"
 
284
            color = "#ff9900"
 
285
        elif (None, None) == (n_rev, m_rev):
 
286
            cluster = None
 
287
            if node in self.common:
 
288
                color = "#6699ff"
 
289
            else:
 
290
                color = "white"
 
291
        elif n_rev is not None:
 
292
            cluster = "my_history"
 
293
            color = "#ffff00"
 
294
        else:
 
295
            assert m_rev is not None
 
296
            cluster = "other_history"
 
297
            color = "#ff0000"
 
298
        if node in self.lcas:
 
299
            color = "#9933cc"
 
300
        if node == self.base:
 
301
            color = "#669933"
 
302
            if node == self.new_base:
 
303
                color = "#33ff33"
 
304
        if node == self.new_base:
 
305
            color = '#33cc99'
 
306
 
 
307
        label = [name]
 
308
        committer, message, nick, date = get_rev_info(node,
 
309
                                                      self.branch.repository)
 
310
        if committer is not None:
 
311
            label.append(committer)
 
312
 
 
313
        if nick is not None:
 
314
            label.append(nick)
 
315
 
 
316
        if date is not None:
 
317
            label.append(date)
 
318
 
 
319
        if node in self.distances:
 
320
            rank = self.distances[node]
 
321
            label.append('d%d' % self.distances[node])
 
322
        else:
 
323
            rank = None
 
324
 
 
325
        d_node = Node("n%d" % num, color=color, label="\\n".join(label),
 
326
                    rev_id=node, cluster=cluster, message=message,
 
327
                    date=date)
 
328
        d_node.rank = rank
 
329
 
 
330
        if node not in self.ancestors:
 
331
            d_node.node_style.append('dotted')
 
332
 
 
333
        return d_node
 
334
 
 
335
    def get_relations(self, collapse=False, max_distance=None):
 
336
        dot_nodes = {}
 
337
        node_relations = []
 
338
        num = 0
 
339
        if collapse:
 
340
            exceptions = self.lcas.union([self.base, self.new_base])
 
341
            visible_ancestors = compact_ancestors(self.descendants,
 
342
                                                  self.ancestors,
 
343
                                                  exceptions)
 
344
        else:
 
345
            visible_ancestors = {}
 
346
            for revision, parents in self.ancestors.iteritems():
 
347
                visible_ancestors[revision] = dict((p, 0) for p in parents)
 
348
        if max_distance is not None:
 
349
            min_distance = max(self.distances.values()) - max_distance
 
350
            visible_ancestors = dict((n, p) for n, p in
 
351
                                     visible_ancestors.iteritems() if
 
352
                                     self.distances[n] >= min_distance)
 
353
        for node, parents in visible_ancestors.iteritems():
 
354
            if node not in dot_nodes:
 
355
                dot_nodes[node] = self.dot_node(node, num)
 
356
                num += 1
 
357
            for parent, skipped in parents.iteritems():
 
358
                if parent not in dot_nodes:
 
359
                    dot_nodes[parent] = self.dot_node(parent, num)
 
360
                    num += 1
 
361
                edge = Edge(dot_nodes[parent], dot_nodes[node])
 
362
                if skipped != 0:
 
363
                    edge.label = "%d" % skipped
 
364
                node_relations.append(edge)
 
365
        return node_relations
 
366
 
 
367
 
 
368
def write_ancestry_file(branch, filename, collapse=True, antialias=True,
 
369
                        merge_branch=None, ranking="forced", max_distance=None):
 
370
    b = Branch.open_containing(branch)[0]
 
371
    if merge_branch is not None:
 
372
        m = Branch.open_containing(merge_branch)[0]
 
373
    else:
 
374
        m = None
 
375
    b.lock_write()
 
376
    try:
 
377
        if m is not None:
 
378
            m.lock_read()
 
379
        try:
 
380
            grapher = Grapher(b, m)
 
381
            relations = grapher.get_relations(collapse, max_distance)
 
382
        finally:
 
383
            if m is not None:
 
384
                m.unlock()
 
385
    finally:
 
386
        b.unlock()
 
387
 
 
388
    ext = filename.split('.')[-1]
 
389
    output = dot_output(relations, ranking)
 
390
    done = False
 
391
    if ext not in RSVG_OUTPUT_TYPES:
 
392
        antialias = False
 
393
    if antialias:
 
394
        output = list(output)
 
395
        try:
 
396
            invoke_dot_aa(output, filename, ext)
 
397
            done = True
 
398
        except NoDot, e:
 
399
            raise BzrCommandError("Can't find 'dot'.  Please ensure Graphviz"\
 
400
                " is installed correctly.")
 
401
        except NoRsvg, e:
 
402
            print "Not antialiasing because rsvg (from librsvg-bin) is not"\
 
403
                " installed."
 
404
            antialias = False
 
405
    if ext in DOT_OUTPUT_TYPES and not antialias and not done:
 
406
        try:
 
407
            invoke_dot(output, filename, ext)
 
408
            done = True
 
409
        except NoDot, e:
 
410
            raise BzrCommandError("Can't find 'dot'.  Please ensure Graphviz"\
 
411
                " is installed correctly.")
 
412
    elif ext == 'dot' and not done:
 
413
        my_file = file(filename, 'wb')
 
414
        for fragment in output:
 
415
            my_file.write(fragment.encode('utf-8'))
 
416
    elif ext == 'html':
 
417
        try:
 
418
            invoke_dot_html(output, filename)
 
419
        except NoDot, e:
 
420
            raise BzrCommandError("Can't find 'dot'.  Please ensure Graphviz"\
 
421
                " is installed correctly.")
 
422
    elif not done:
 
423
        print "Unknown file extension: %s" % ext