69f392b30212b1bc842c95ae3527eb2aff586dc2
[hooke.git] / hooke / util / graph.py
1 # Copyright (C) 2010-2012 W. Trevor King <wking@tremily.us>
2 #
3 # This file is part of Hooke.
4 #
5 # Hooke is free software: you can redistribute it and/or modify it under the
6 # terms of the GNU Lesser General Public License as published by the Free
7 # Software Foundation, either version 3 of the License, or (at your option) any
8 # later version.
9 #
10 # Hooke is distributed in the hope that it will be useful, but WITHOUT ANY
11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
12 # A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
13 # details.
14 #
15 # You should have received a copy of the GNU Lesser General Public License
16 # along with Hooke.  If not, see <http://www.gnu.org/licenses/>.
17
18 """Define :class:`Graph`, a directed, acyclic graph structure.
19 :class:`Graph`\s are composed of :class:`Node`\s, also defined by this
20 module.
21 """
22
23 import bisect
24
25 class CyclicGraphError (ValueError):
26     pass
27
28 class Node (list):
29     """A node/element in a graph.
30
31     Contains a list of the node's parents, and stores the node's
32     `data`.
33
34     Examples
35     --------
36
37     >>> a = Node(data='a')
38     >>> b = Node(parents=[a], data='b')
39     >>> c = Node(parents=[a], data='c')
40     >>> d = Node(parents=[b, c], data='d')
41     >>> str(d)
42     'd'
43
44     We can list all of a node's ancestors.
45
46     >>> print [node for node in d.ancestors()]
47     [b, c, a]
48     >>> print [node for node in d.ancestors(depth_first=True)]
49     [b, a, c]
50
51     Ancestors works with cycles.
52
53     >>> a.append(d)
54     >>> print [node for node in d.ancestors()]
55     [b, c, a, d]
56
57     We can find the cycle path.
58
59     >>> print d.parent_path(d)
60     [b, a, d]
61
62     After a run through :meth:`Graph.set_children`, we can also
63     list children
64
65     >>> g = Graph([a, b, c, d])
66     >>> g.set_children()
67     >>> print a.children
68     [b, c]
69
70     And descendents.
71
72     >>> print [node for node in a.descendents(depth_first=True)]
73     [b, d, a, c]
74     """
75     def __init__(self, parents=[], data=None):
76         list.__init__(self, parents)
77         self.data = data
78         self.children = []
79
80     def __cmp__(self, other):
81         return -cmp(self.data, other.data)
82     def __eq__(self, other):
83         return self.__cmp__(other) == 0
84     def __ne__(self, other):
85         return self.__cmp__(other) != 0
86     def __lt__(self, other):
87         return self.__cmp__(other) < 0
88     def __gt__(self, other):
89         return self.__cmp__(other) > 0
90     def __le__(self, other):
91         return self.__cmp__(other) <= 0
92     def __ge__(self, other):
93         return self.__cmp__(other) >= 0
94
95     def __str__(self):
96         return str(self.data)
97     def __repr__(self):
98         return self.__str__()
99
100     def traverse(self, next, depth_first=False):
101         """Iterate through all nodes returned by `next(node)`.
102
103         Will only yield each traversed node once, even in the case of
104         diamond inheritance, etc.
105
106         Breadth first by default.  Set `depth_first==True` for a
107         depth first search.
108         """
109         stack = list(next(self))
110         popped = []
111         while len(stack) > 0:
112             node = stack.pop(0)
113             if node in popped:
114                 continue
115             popped.append(node)
116             yield node
117             if depth_first == True:
118                 for target in reversed(next(node)):
119                     stack.insert(0, target)
120             else:
121                 stack.extend(next(node))
122
123     def ancestors(self, depth_first=False):
124         """Generate all ancestors.
125
126         This is a small wrapper around :meth:`traverse`.
127         """
128         next = lambda node : node  # list node's parents
129         for node in self.traverse(next=next, depth_first=depth_first):
130             yield node
131
132     def descendents(self, depth_first=False):
133         """Generate all decendents.
134
135         This is a small wrapper around :meth:`traverse`.
136         """
137         next = lambda node : node.children
138         for node in self.traverse(next=next, depth_first=depth_first):
139             yield node
140
141     def path(self, next, node):
142         """Return the shortest list of nodes connecting `self` to
143         `node` via `next(node)`.
144         """
145         if node in self:
146             return [node]
147         stack = list(next(self))
148         paths = dict((id(n), [n]) for n in stack)
149         while len(stack) > 0:
150             n = stack.pop(0)
151             n_path = paths[id(n)]
152             for target in next(n):
153                 if id(target) in paths:
154                     continue
155                 t_path = list(n_path)
156                 t_path.append(target)
157                 if id(target) == id(node):
158                     return t_path
159                 stack.append(target)
160                 paths[id(target)] = t_path
161
162     def parent_path(self, node):
163         """Return the shortest list of nodes connecting `self` to
164         its parent `node`.
165
166         This is a small wrapper around :meth:`path`.
167         """
168         next = lambda node : node  # list node's parents
169         return self.path(next, node)
170
171     def child_path(self, node):
172         """Return the shortest list of nodes connecting `self` to
173         its child `node`.
174
175         This is a small wrapper around :meth:`path`.
176         """
177         next = lambda node : node.children
178         return self.path(next, node)
179
180
181 class GraphRow (object):
182     """Represent the state of a single row in a graph.
183
184     Generated by :class:`GraphRowGenerator`, printed with
185     :class:`GraphRowPrinter`.
186
187     :attr:`node` is the active node and :attr:`active` is its branch
188     column index.  :attr:`width` is the number of current branch
189     columns.
190
191     :attr:`born`, :attr:`dead`, and :attr:`inherited` are lists of
192     `(branch_column_index, target_node)` pairs.  `dead` lists nodes
193     from previous rows whose branches complete on this row,
194     `inherited` lists nodes from previous rows whose branches continue
195     through this row, and `born` list nodes whose branches start on
196     this row.
197     """
198     def __init__(self, node, active=-1, dead=None, inherited=None, born=None,
199                  tip_to_root=False):
200         self.node = node
201         self.active = active
202         if dead == None:
203             dead = []
204         self.dead = dead
205         if inherited == None:
206             inherited = []
207         self.inherited = inherited
208         if born == None:
209             born = []
210         self.born = born
211         self.tip_to_root = tip_to_root
212
213 class GraphRowPrinter (object):
214     """Customizable printing for :class:`GraphRow`.
215
216     The string rendering can be customized by changing :attr:`chars`.
217     Control over the branch columns:
218
219     ================= ===========================================
220     `node: ...`       the active (most recently inserted) node
221     `split/join: ...` branching/merging runs from the active node
222     `run: connected`  connect a branch to its eventual node
223     `run: blank`      place-holder for extinct runs
224     ================= ===========================================
225
226     Branch columns are seperated by separator columns:
227
228     ================= =======================================================
229     `sep: split/join` separate split/join branch columns from the active node
230     `sep: default`    separate all remaining branch columns
231     ================= =======================================================
232     """
233     def __init__(self, chars=None):
234         if chars == None:
235             chars = {
236                 'node: both tip and root': 'b',
237                 'node: root': 'r',
238                 'node: tip': 't',
239                 'node: regular': '*',
240                 'split/join: born and died left of active': '>',
241                 'split/join: born and died right of active': '<',
242                 'split/join: born left of active': '/',
243                 'split/join: born right of active': '\\',
244                 'split/join: died left of active': '\\',
245                 'split/join: died right of active': '/',
246                 'run: blank': ' ',
247                 'run: connected': '|',
248                 'sep: split/join': '-',
249                 'sep: default': ' ',               
250                 }
251         self.chars = chars
252     def __call__(self, graph_row):
253         """Render the :class:`GraphRow` instance `graph_row` as a
254         string.
255         """
256         dead = [i for i,node in graph_row.dead]
257         inherited = [i for i,node in graph_row.inherited]
258         born = [i for i,node in graph_row.born]
259         right_connect = max(graph_row.active,
260                             max(born+[-1]), # +[-1] protects against empty born
261                             max(dead+[-1]))
262         left_connect = min(graph_row.active,
263                            min(born+[right_connect]),
264                            min(dead+[right_connect]))
265         max_col = max(right_connect, max(inherited+[-1]))
266         string = []
267         for i in range(max_col + 1):
268             # Get char, the node or branch column character.
269             if i == graph_row.active:
270                 if len(born) == 0:
271                     if len(dead) == 0:
272                         char = self.chars['node: both tip and root']
273                     elif graph_row.tip_to_root == True:
274                         # The dead are children
275                         char = self.chars['node: root']
276                     else: # The dead are parents
277                         char = self.chars['node: tip']
278                 elif len(dead) == 0:
279                     if graph_row.tip_to_root == True:
280                         # The born are parents
281                         char = self.chars['node: tip']
282                     else: # The born are children
283                         char = self.chars['node: root']
284                 else:
285                     char = self.chars['node: regular']
286             elif i in born:
287                 if i in dead: # born and died
288                     if i < graph_row.active:
289                         char = self.chars[
290                             'split/join: born and died left of active']
291                     else:
292                         char = self.chars[
293                             'split/join: born and died right of active']
294                 else: # just born
295                     if i < graph_row.active:
296                         char = self.chars['split/join: born left of active']
297                     else:
298                         char = self.chars['split/join: born right of active']
299             elif i in dead: # just died
300                 if i < graph_row.active:
301                     char = self.chars['split/join: died left of active']
302                 else:
303                     char = self.chars['split/join: died right of active']
304             elif i in inherited:
305                 char = self.chars['run: connected']
306             else:
307                 char = self.chars['run: blank']
308             # Get sep, the separation character.
309             if i < left_connect or i >= right_connect:
310                 sep = self.chars['sep: default']
311             else:
312                 sep = self.chars['sep: split/join']
313                 if char == self.chars['run: blank']:
314                     char = self.chars['sep: split/join']
315             string.extend([char, sep])
316         return ''.join(string)[:-1] # [-1] strips final sep
317
318 class GraphRowGenerator (list):
319     """A :class:`GraphRow` generator.
320
321     Contains a list of :class:`GraphRow`\s (added with
322     :meth:`insert`(:class:`hooke.util.graph.Node`)).  You should
323     generate a graph with repeated calls::
324
325         tip_to_root = True
326         g = GraphRowGenerator(tip_to_root=tip_to_root)
327         p = GraphRowPrinter(tip_to_root=tip_to_root)
328         for node in nodes:
329             g.insert(node)
330             print p(g[-1])
331
332     For the split/join branch columns, "born" and "dead" are defined
333     from the point of view of `GraphRow`.  For root-to-tip ordering
334     (`tip_to_root==False`, the default), "born" runs are determined
335     by the active node's children (which have yet to be printed) and
336     "dead" runs by its parents (which have been printed).  If
337     `tip_to_root==True`, swap "children" and "parents" in the above
338     sentence.
339     """
340     def __init__(self, tip_to_root=False):
341         list.__init__(self)
342         self.tip_to_root = tip_to_root
343     def insert(self, node):
344         """Add a new node to the graph.
345
346         If `tip_to_root==True`, nodes should be inserted in
347         tip-to-root topological order (i.e. node must be inserted
348         before any of its parents).
349
350         If `tip_to_root==False`, nodes must be inserted before any
351         of their children.
352         """
353         if len(self) == 0:
354             previous = GraphRow(node=None, active=-1)
355         else:
356             previous = self[-1]
357         current = GraphRow(node=node, active=-1, tip_to_root=self.tip_to_root)
358         if self.tip_to_root == True: # children die, parents born
359             dead_nodes = list(current.node.children)
360             born_nodes = list(current.node)
361         else: # root-to-tip: parents die, children born
362             dead_nodes = list(current.node)
363             born_nodes = list(current.node.children)
364         # Mark the dead and inherited branch columns
365         for i,node in previous.inherited + previous.born:
366             if node in dead_nodes or node == current.node:
367                 current.dead.append((i, node))
368             else:
369                 current.inherited.append((i, node))
370         # Place born and active branch columns
371         num_born = max(len(born_nodes), 1) # 1 to ensure slot for active node
372         remaining = num_born # number of nodes left to place
373         used_slots = [i for i,n in current.inherited]
374         old_max = max(used_slots+[-1]) # +[-1] in case used_slots is empty
375         slots = sorted([i for i in range(old_max+1) if i not in used_slots])
376         remaining -= len(slots)
377         slots.extend(range(old_max+1, old_max+1+remaining))
378         current.active = slots[0]
379         current.born = zip(slots, born_nodes)
380         # TODO: sharing branches vs. current 1 per child
381         self.append(current)
382
383
384 class Graph (list):
385     """A directed, acyclic graph structure.
386
387     Contains methods for sorting and printing graphs.
388
389     Examples
390     --------
391
392     >>> class Nodes (object): pass
393     >>> n = Nodes()
394     >>> for char in ['a','b','c','d','e','f','g','h','i']:
395     ...     setattr(n, char, Node(data=char))
396     >>> n.b.append(n.a)
397     >>> n.c.append(n.a)
398     >>> n.d.append(n.a)
399     >>> n.e.extend([n.b, n.c, n.d])
400     >>> n.f.append(n.e)
401     >>> n.g.append(n.e)
402     >>> n.h.append(n.e)
403     >>> n.i.extend([n.f, n.g, n.h])
404     >>> g = Graph([n.a,n.b,n.c,n.d,n.e,n.f,n.g,n.h,n.i])
405     >>> g.topological_sort(tip_to_root=True)
406     >>> print [node for node in g]
407     [i, h, g, f, e, d, c, b, a]
408     >>> print g.ascii_graph()
409     r-\-\ a
410     | | * b
411     | * | c
412     * | | d
413     *-<-< e
414     | | * f
415     | * | g
416     * | | h
417     t-/-/ i
418     >>> print g.ascii_graph(tip_to_root=True)
419     t-\-\ i
420     | | * h
421     | * | g
422     * | | f
423     *-<-< e
424     | | * d
425     | * | c
426     * | | b
427     r-/-/ a
428
429     >>> for char in ['a','b','c','d','e','f','g','h']:
430     ...     setattr(n, char, Node(data=char))
431     >>> n.b.append(n.a)
432     >>> n.c.append(n.b)
433     >>> n.d.append(n.a)
434     >>> n.e.append(n.d)
435     >>> n.f.extend([n.b, n.d])
436     >>> n.g.extend([n.e, n.f])
437     >>> n.h.extend([n.c, n.g])
438     >>> g = Graph([n.a,n.b,n.c,n.d,n.e,n.f,n.g,n.h])
439     >>> print g.ascii_graph(tip_to_root=True)
440     t-\ h
441     | *-\ g
442     | | *-\ f
443     | * | | e
444     | *-|-/ d
445     * | | c
446     *-|-/ b
447     r-/ a
448
449     >>> for char in ['a','b','c','d','e','f','g','h','i']:
450     ...     setattr(n, char, Node(data=char))
451     >>> for char in ['a', 'b','c','d','e','f','g','h']:
452     ...     nx = getattr(n, char)
453     ...     n.i.append(nx)
454     >>> g = Graph([n.a,n.b,n.c,n.d,n.e,n.f,n.g,n.h,n.i])
455     >>> print g.ascii_graph(tip_to_root=True)
456     t-\-\-\-\-\-\-\ i
457     | | | | | | | r h
458     | | | | | | r g
459     | | | | | r f
460     | | | | r e
461     | | | r d
462     | | r c
463     | r b
464     r a
465
466     >>> for char in ['a','b','c','d','e','f','g','h','i']:
467     ...     setattr(n, char, Node(data=char))
468     >>> for char in ['b','c','d','e','f','g','h','i']:
469     ...     nx = getattr(n, char)
470     ...     nx.append(n.a)
471     >>> g = Graph([n.a,n.b,n.c,n.d,n.e,n.f,n.g,n.h,n.i])
472     >>> print g.ascii_graph(tip_to_root=True)
473     t i
474     | t h
475     | | t g
476     | | | t f
477     | | | | t e
478     | | | | | t d
479     | | | | | | t c
480     | | | | | | | t b
481     r-/-/-/-/-/-/-/ a
482
483     >>> for char in ['a','b','c','d','e','f','g','h','i']:
484     ...     setattr(n, char, Node(data=char))
485     >>> n.d.append(n.a)
486     >>> n.e.extend([n.a, n.c])
487     >>> n.f.extend([n.c, n.d, n.e])
488     >>> n.g.extend([n.b, n.e, n.f])
489     >>> n.h.extend([n.a, n.c, n.d, n.g])
490     >>> n.i.extend([n.a, n.b, n.c, n.g])
491     >>> g = Graph([n.a,n.b,n.c,n.d,n.e,n.f,n.g,n.h,n.i])
492     >>> print g.ascii_graph(tip_to_root=True)
493     t-\-\-\ i
494     | | | | t-\-\-\ h
495     | | | *-|-|-|-<-\ g
496     | | | | | | | | *-\-\ f
497     | | | | | | | *-|-|-< e
498     | | | | | | *-|-|-/ | d
499     | | r-|-|-/-|-|-/---/ c
500     | r---/ |   | | b
501     r-------/---/-/ a
502
503     Ok, enough pretty graphs ;).  Here's an example of cycle
504     detection.
505
506     >>> for char in ['a','b','c','d']:
507     ...     setattr(n, char, Node(data=char))
508     >>> n.b.append(n.a)
509     >>> n.c.append(n.a)
510     >>> n.d.extend([n.b, n.c])
511     >>> n.a.append(n.d)
512     >>> g = Graph([n.a,n.b,n.c,n.d])
513     >>> g.check_for_cycles()
514     Traceback (most recent call last):
515       ...
516     CyclicGraphError: cycle detected:
517       a
518       d
519       b
520       a
521     """
522     def set_children(self):
523         """Fill out each node's :attr:`Node.children` list.
524         """
525         for node in self:
526             for parent in node:
527                 if node not in parent.children:
528                     parent.children.append(node)
529
530     def topological_sort(self, tip_to_root=False):
531         """Algorithm from git's commit.c `sort_in_topological_order`_.
532
533         Default ordering is root-to-tip.  Set `tip_to_root=True` for
534         tip-to-root.
535
536         In situations where topological sorting is ambiguous, the
537         nodes are sorted using the node comparison functions (__cmp__,
538         __lt__, ...).  If `tip_to_root==True`, the inverse
539         comparison functions are used.
540
541         .. _sort_in_topological_order:
542           http://git.kernel.org/?p=git/git.git;a=blob;f=commit.c;h=731191e63bd39a89a8ea4ed0390c49d5605cdbed;hb=HEAD#l425
543         """
544         # sort tip-to-root first, then reverse if neccessary
545         for node in self:
546             node._outcount = 0
547         for node in self:
548             for parent in node:
549                 parent._outcount += 1
550         tips = sorted([node for node in self if node._outcount == 0])
551         orig_len = len(self)
552         del self[:]
553         while len(tips) > 0:
554             node = tips.pop(0)
555             for parent in node:
556                 parent._outcount -= 1
557                 if parent._outcount == 0:
558                     bisect.insort(tips, parent)
559             node._outcount = -1
560             self.append(node)
561         final_len = len(self)
562         if final_len != orig_len:
563             raise CyclicGraphError(
564                 '%d of %d elements not reachable from tips'
565                 % (orig_len - final_len, orig_len))
566         if tip_to_root == False:
567             self.reverse()
568
569     def check_for_cycles(self):
570         """Check for cyclic references.
571         """
572         for node in self:
573             if node in node.ancestors():
574                 path = node.parent_path(node)
575                 raise CyclicGraphError(
576                     'cycle detected:\n  %s'
577                     % '\n  '.join([repr(node)]+[repr(node) for node in path]))
578
579     def graph_rows(self, tip_to_root=False):
580         """Generate a sequence of (`graph_row`, `node`) tuples.
581
582         Preforms :meth:`set_children` and :meth:`topological_sort`
583         internally.
584         """
585         graph_row_generator = GraphRowGenerator(tip_to_root=tip_to_root)
586         self.set_children()
587         self.topological_sort(tip_to_root=tip_to_root)
588         for node in self:
589             graph_row_generator.insert(node)
590             yield (graph_row_generator[-1], node)
591
592     def ascii_graph(self, graph_row_printer=None, string_fn=str,
593                     tip_to_root=False):
594         """Print an ascii graph on the left with `string_fn(node)` on
595         the right.  If `graph_row_printer` is `None`, a default
596         instance of :class:`GraphRowPrinter` will be used.
597
598         See the class docstring for example output.
599         """
600         if graph_row_printer == None:
601             graph_row_printer = GraphRowPrinter()
602         graph = []
603         for row,node in self.graph_rows(tip_to_root=tip_to_root):
604             graph.append('%s %s' % (graph_row_printer(row), string_fn(node)))
605         return '\n'.join(graph)