planetstack/ec2_observer/toposort.py

   1 #!/usr/bin/python
   2
   3 import time
   4 import traceback
   5 import commands
   6 import threading
   7 import json
   8 import pdb
   9
  10 from datetime import datetime
  11 from collections import defaultdict
  12
  13 # Assumes that there are no empty dependencies
  14 # in the graph. E.g. Foo -> []
  15 def dfs(graph, visit):
  16         nodes = graph.keys()
  17         edge_nodes = set()
  18
  19         for n in nodes:
  20                 edge_nodes|=set(graph[n])
  21
  22         sinks = list(edge_nodes - set(nodes))
  23         sources = list(set(nodes) - edge_nodes)
  24
  25         nodes.extend(sinks)
  26
  27         visited = set(sources)
  28         stack = sources
  29         while stack:
  30                 current = stack.pop()
  31                 visit(current)
  32                 for node in graph[current]:
  33                         if node not in visited:
  34                                 stack.append(node)
  35                                 visited.add(node)
  36
  37         return sources
  38
  39 # Topological sort
  40 # Notes:
  41 # - Uses a stack instead of recursion
  42 # - Forfeits optimization involving tracking currently visited nodes
  43 def toposort(g, steps=None):
  44         # Get set of all nodes, including those without outgoing edges
  45         keys = set(g.keys())
  46         values = set({})
  47         for v in g.values():
  48                 values=values | set(v)
  49
  50         all_nodes=list(keys|values)
  51         if (not steps):
  52                 steps = all_nodes
  53
  54         # Final order
  55         order = []
  56
  57         # DFS stack, not using recursion
  58         stack = []
  59
  60         # Unmarked set
  61         unmarked = all_nodes
  62
  63         # visiting = [] - skip, don't expect 1000s of nodes, |E|/|V| is small
  64
  65         while unmarked:
  66                 stack.insert(0,unmarked[0]) # push first unmarked
  67
  68                 while (stack):
  69                         n = stack[0]
  70                         add = True
  71                         try:
  72                                 for m in g[n]:
  73                                         if (m in unmarked):
  74                                                 if (m not in stack):
  75                                                         add = False
  76                                                         stack.insert(0,m)
  77                                                 else:
  78                                                         # Should not happen, if so there's a loop
  79                                                         print 'Loop at %s'%m
  80                         except KeyError:
  81                                 pass
  82                         if (add):
  83                                 if (n in steps):
  84                                         order.append(n)
  85                                 item = stack.pop(0)
  86                                 unmarked.remove(item)
  87
  88         noorder = list(set(steps) - set(order))
  89         return order + noorder
  90
  91 def main():
  92         graph_file=open('planetstack.deps').read()
  93         g = json.loads(graph_file)
  94         print toposort(g)
  95
  96 if (__name__=='__main__'):
  97         main()
  98
  99 #print toposort({'a':'b','b':'c','c':'d','d':'c'},['d','c','b','a'])