1 files changed, 28 insertions, 22 deletions
diff --git a/kpov_judge/web/kpov_judge/kpov_draw_setup.py b/kpov_judge/web/kpov_judge/kpov_draw_setup.py
index a6ffc82..0042389 100644
--- a/kpov_judge/web/kpov_judge/kpov_draw_setup.py
+++ b/kpov_judge/web/kpov_judge/kpov_draw_setup.py
@@ -1,35 +1,41 @@
 import pygraphviz as pgv
 
-def draw_setup(computers, networks, destination = None):
-    G = pgv.AGraph()
+def draw_setup(computers, networks, destination=None,
+               icon_prefix = '../../../static/icons/',
+               format='svg', icon_path = '', icon_suffix = None):
+    if icon_suffix is None:
+        icon_suffix = format
+    icon_suffix = '.' + icon_suffix
+    G = pgv.AGraph(imagepath=icon_path + '/')
+    print G.graph_attr
     have_internet = []
-    for net, net_params in networks.iteritems():
-        if net_params.get('public', False):
-            have_internet.append(net)
-        G.add_node('net-' + net, label=net, shape='rectangle')
+    for net in networks:
+        net_name = net.get('name', 'net')
+        if net.get('public', False):
+            have_internet.append(net_name)
+        G.add_node('net-' + net_name, label=net_name, shape='rectangle')
     if len(have_internet):
         G.add_node('net-' + 'internet', 
                    label='internet',
-                   labelloc='b', image='icons/internet.svg', shape='none')
+                   labelloc='b', 
+                   image=icon_prefix + 'internet' + icon_suffix, 
+                   shape='none')
         for n in have_internet:
             G.add_edge('net-' + n, 'net-internet')
-    for c, properties in computers.iteritems():
-        G.add_node('comp-' + c,
-                   label = c,
-                   image='icons/computer.svg', 
-                   shape='none', labelloc='b')
-        G_sub = G.add_subgraph(['comp-' + c], name='cluster-' + c,
-                               label = c,
-                               rank='same', pack=True)
+    for properties in computers:
+        c = properties.get('name', '')
+        label_str = '< <TABLE BORDER="0"><TR><TD COLSPAN="2"><B>{}</B></TD></TR><TR><TD COLSPAN="2"><IMG SRC="{}"/></TD></TR>'
+        label = label_str.format(c, icon_prefix + 'computer' + icon_suffix)
         for hdd in properties.get('disks', []):
-            G_sub.add_node('disk-' + hdd['name'],
-                       label=hdd['name'],
-                       shape='none',
-                       image='icons/drive-harddisk.svg')
-            G_sub.add_edge('comp-' + c, 'disk-' + hdd['name'])
-        for iface in properties.get('network_interfaces'):
+            icon = icon_prefix + 'drive-harddisk' + icon_suffix
+            label += '<TR><TD><IMG SRC="{}" /></TD><TD>{}</TD></TR>'.format(icon, hdd['name'])
+        label += '</TABLE> >'
+        G.add_node('comp-' + c,
+                   label = label,
+                   shape='box', labelloc='b')
+        for iface in properties.get('network_interfaces', []):
             G.add_edge('comp-' + c, 'net-' + iface['network'])
-    return G.draw(path=destination, format='svg', prog='dot')
+    return G.draw(path=destination, format=format, prog='dot')
 
 if __name__ == '__main__':
     import sample_task as task