bin/tools/graph.py

   1 #!/usr/bin/python
   2
   3 class graph(object):
   4         def __init__(self, nodes, transitions):
   5                 self.nodes = nodes
   6                 trans = {}
   7                 for t in transitions:
   8                         trans.setdefault(t[0], [])
   9                         trans[t[0]].append(t[1])
  10                 self.transitions = trans
  11                 self.result = {}
  12                 self.levels = {}
  13
  14         def process_ranking(self, node, level=0):
  15                 if node not in self.result:
  16                         self.result[node] = {'x': None, 'y':level}
  17                 else:
  18                         if level > self.result[node]['y']:
  19                                 self.result[node]['y'] = level
  20                 for t in self.transitions.get(node, []):
  21                         self.process_ranking(t, level+1)
  22
  23         def preprocess_order(self):
  24                 levels = {}
  25                 for r in self.result:
  26                         l = self.result[r]['y']
  27                         levels.setdefault(l,[])
  28                         levels[l].append(r)
  29                 self.levels = levels
  30
  31         def process_order(self, level):
  32                 self.levels[level].sort(lambda x,y: cmp(self.result[x]['x'], self.result[y]['x']))
  33                 for nodepos in range(len(self.levels[level])):
  34                         node = self.levels[level][nodepos]
  35                         if nodepos == 0:
  36                                 left = self.result[node]['x']- 0.5
  37                         else:
  38                                 left = (self.result[node]['x'] + self.result[self.levels[level][nodepos-1]]['x']) / 2.0
  39
  40                         if nodepos == (len(self.levels[level])-1):
  41                                 right = self.result[node]['x'] + 0.5
  42                         else:
  43                                 right = (self.result[node]['x'] +  self.result[self.levels[level][nodepos+1]]['x']) / 2.0
  44
  45
  46                         if self.transitions.get(node, False):
  47                                 if len(self.transitions[node])==1:
  48                                         pos = (left+right)/2.0
  49                                         step = 0
  50                                 else:
  51                                         pos = left
  52                                         step = (-left+right) / (len(self.transitions[node])-1)
  53
  54                                 for n2 in self.transitions[node]:
  55                                         self.result[n2]['x'] = pos
  56                                         pos += step
  57
  58         def process(self, starting_node):
  59                 pos = (len(starting_node) - 1.0)/2.0
  60                 for s in starting_node:
  61                         self.process_ranking(s)
  62                         self.result[s]['x'] = pos
  63                         pos += 1.0
  64
  65                 self.preprocess_order()
  66
  67                 for n in self.levels:
  68                         self.process_order(n)
  69
  70         def __str__(self):
  71                 result = ''
  72                 for l in self.levels:
  73                         result += 'PosY: ' + str(l) + '\n'
  74                         for node in self.levels[l]:
  75                                 result += '\tPosX: '+ str(self.result[node]['x']) + '  - Node:' + node + "\n"
  76                 return result
  77
  78         def scale(self, maxx, maxy, plusx2=0, plusy2=0):
  79                 plusx = - min(map(lambda x: x['x'],self.result.values()))
  80                 plusy = - min(map(lambda x: x['y'],self.result.values()))
  81
  82                 maxcurrent = 1.0
  83                 for l in self.levels:
  84                         for n in range(1, len(self.levels[l])):
  85                                 n1 = self.levels[l][n]
  86                                 n2 = self.levels[l][n-1]
  87                                 diff = abs(self.result[n2]['x']-self.result[n1]['x'])
  88                                 if diff<maxcurrent:
  89                                         maxcurrent=diff
  90                 factor = maxx / diff
  91                 for r in self.result:
  92                         self.result[r]['x'] = (self.result[r]['x']+plusx) * factor + plusx2
  93                         self.result[r]['y'] = (self.result[r]['y']+plusy) * maxy + plusy2
  94
  95         def result_get(self):
  96                 return self.result
  97
  98 if __name__=='__main__':
  99         starting_node = ['mrp']  # put here nodes with flow_start=True
 100         nodes = ['project','account','hr','base','product','mrp','test']
 101         transitions = [
 102                 ('mrp','project'),
 103                 ('project','product'),
 104                 ('project','account'),
 105                 ('project','hr'),
 106                 ('product','base'),
 107                 ('account','product'),
 108                 ('account','test'),
 109                 ('account','base'),
 110                 ('hr','base'),
 111                 ('test','base')
 112         ]
 113
 114         radius = 20
 115         g = graph(nodes, transitions)
 116         g.process(starting_node)
 117         g.scale(radius*3,radius*3, radius, radius)
 118
 119         print g
 120
 121         import Image
 122         import ImageDraw
 123         img = Image.new("RGB", (800, 600), "#ffffff")
 124         draw = ImageDraw.Draw(img)
 125
 126         for name,node in g.result.items():
 127                 draw.arc( (int(node['x']-radius), int(node['y']-radius),int(node['x']+radius), int(node['y']+radius) ), 0, 360, (128,128,128))
 128                 draw.text( (int(node['x']),  int(node['y'])), name,  (128,128,128))
 129
 130
 131         for nodefrom in g.transitions:
 132                 for nodeto in g.transitions[nodefrom]:
 133                         draw.line( (int(g.result[nodefrom]['x']), int(g.result[nodefrom]['y']),int(g.result[nodeto]['x']),int(g.result[nodeto]['y'])),(128,128,128) )
 134
 135         img.save("graph.png", "PNG")
 136