dependency graphs (using heapq)

gistfile1.py

# dependency graphs.
# utilities to declare dependencies between items and get which ones to do first.
# this is using heapq for walking up a dependency chain.

from heapq import *

class DependencyGraph(object):
   '''
   Represent a dependency graph.
   All nodes should have dependencies presented as a list of node names.

   graph = DependencyGraph({
      'B.1': ['A.0'],
      'C.2': ['B.1'],
   })

   '''

   class CyclicError(Exception):
      '''
      error raised when a cycle is detected
      '''
      pass

   # default node representation
   # we will use a heapq to navigate the nodes from the bottom to the top.
   # small priority: requirement/dependency, large priority: parent/aggregate
   defaultnode = {'priority': 0}

   def __init__(self, data=None):
      # parent_nodes and children_nodes are to store the edges in the graph
      self.parent_nodes = {}
      self.children_nodes = {}
      # nodes is there to store the data associated with a node (e.g. priority)
      self.nodes = {}
      # quick access to nodes with no children
      self.leaves = set()
      if data is not None:
         self.insert_graph(data)

   def walk_priority(self, start=None):
      '''
      iterate over all nodes that depend on a starting set.
      we use the heapq to navigate the nodes from the bottom to the top.
      '''
      nodes = []
      visited = set()
      if start is None:
         start = self.leaves
      for target in start:
         priority = self.nodes.get(target, self.defaultnode)['priority']
         heappush(nodes, (priority, target))
         visited.add(target)
      while nodes:
         priority, target = heappop(nodes)
         yield (priority, target)
         for parent in self.parent_nodes[target]:
            if parent in visited:
               continue
            node = self.nodes[parent]
            parentpriority = node['priority']
            heappush(nodes, (parentpriority, parent))
            visited.add(parent)

   def fix_priority(self, target):
      '''
      fix priorities in the top of the graph.
      since nodes can be inserted out of order, we need to recalculate
      priorities for anything above the current target.
      '''
      nodes = []
      priority = self.nodes[target]['priority']
      origpriority, origtarget = priority, target
      heappush(nodes, (priority, target))
      while nodes:
         priority, target = heappop(nodes)
         if target == origtarget and priority > origpriority:
            raise self.CyclicError('cannot insert %s' % target)
         # expected parent priority is at least 1 more
         parentpriority = priority + 1
         for parent in self.parent_nodes[target]:
            node = self.nodes[parent]
            if node['priority'] < parentpriority:
               node['priority'] = parentpriority
               # parent changed, queue it for processing
               heappush(nodes, (parentpriority, parent))

   def insert_node(self, target, dependencies=None):
      '''
      insert a node out of order.
      '''
      # defaults
      children = self.children_nodes.setdefault(target, set())
      parents = self.parent_nodes.setdefault(target, set())
      self.nodes.setdefault(target, self.defaultnode.copy())
      if dependencies:
         # fix target children
         children = self.children_nodes.setdefault(target, set())
         children.update(dependencies)
         # fix children parents & get priority
         childpriority = 0
         for child in dependencies:
            childparents = self.parent_nodes.setdefault(child, set())
            childparents.add(target)
            # excuse child not yet inserted.
            childnode = self.nodes.get(child, self.defaultnode)
            childpriority = max(childpriority, childnode['priority'])
         self.nodes[target]['priority'] = childpriority + 1
      else:
         self.leaves.add(target)
      # fix ascendants priority
      self.fix_priority(target)

   def insert_graph(self, data):
      '''
      insert nodes out of order from data items.
      '''
      for target, dependencies in data.items():
         self.insert_node(target, dependencies)



def test1():
   '''
   test general graph structure

    5:            A
                  |
                  |
    4:            B
                 /|\
                / | \ 
    3:         /  |  D
              |   |   \
              |   |    \
    2:        C   |     I
              |   |    /|\
              |   |   / | \
    1:        E   |  K  |  \
            / | \ | /   |   \
           /  |  \|/    |    \
    0:    F   G*  H     G*    J

   '''
   graph = DependencyGraph({
      'A.5': ['B.4'],
      'B.4': ['C.2', 'H.0', 'D.3'],
      'C.2': ['E.1'],
      'D.3': ['I.2'],
      'E.1': ['F.0', 'G.0', 'H.0'],
      'F.0': None,
      'G.0': None,
      'H.0': None,
      'I.2': ['G.0', 'K.1', 'J.0'],
      'J.0': None,
      'K.1': ['H.0'],
   })
   for priority, target in graph.walk_priority():
      print 'priority:', priority, 'name:', target
   for priority, target in graph.walk_priority():
      assert target.endswith('.%d' % priority), "wrong priority detected"
   print 'test successful'


def test2():
   '''
   test cyclic errors.

      0        1        2      [err]
     A* <----- B <----- C <----- A*

   '''
   graph = DependencyGraph({
      'B.1': ['A.0'],
      'C.2': ['B.1'],
   })
   for priority, target in graph.walk_priority(start=['A.0']):
      print 'priority:', priority, 'name:', target
   try:
      graph.insert_node('A.0', ['C.2'])
      print 'test failed'
   except graph.CyclicError:
      print 'test successful'