samba-tool visualize: group (and colour) DCs by site

[kai/samba-autobuild/.git] / python / samba / graph.py

diff --git a/python/samba/graph.py b/python/samba/graph.py
index f626287800d111c943f56da9c538dc49a93c0f64..7d195e3574a75d81f377714109f9a8c4d203f79b 100644 (file)
--- a/python/samba/graph.py
+++ b/python/samba/graph.py
@@ -19,8 +19,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from __future__ import print_function
+from __future__ import division
 from samba import colour
 import sys
+from itertools import cycle, groupby
 
 FONT_SIZE = 10
 
@@ -39,8 +41,8 @@ def reformat_graph_label(s):
             if '-' in p[2:20]:
                 q, p = p.split('-', 1)
             else:
-                n = len(p) / 12
-                b = len(p) / n
+                n = len(p) // 12
+                b = len(p) // n
                 q, p = p[:b], p[b:]
             pieces.append(q + '-')
         if p:
@@ -129,7 +131,7 @@ def shorten_vertex_names(edges, vertices, suffix=',...', aggressive=False):
                 break
         else:
             map = dict((k, v.replace(',CN=Servers,', ',**,'))
-                       for k, v in map.iteritems())
+                       for k, v in map.items())
             replacements.append(('**', 'CN=Servers'))
 
         for v in vertices2:
@@ -137,7 +139,7 @@ def shorten_vertex_names(edges, vertices, suffix=',...', aggressive=False):
                 break
         else:
             map = dict((k, v.replace('CN=NTDS Settings,', '*,'))
-                       for k, v in map.iteritems())
+                       for k, v in map.items())
             replacements.append(('*', 'CN=NTDS Settings'))
 
         edges2 = [(map.get(a, a), map.get(b, b)) for a, b in edges2]
@@ -463,11 +465,11 @@ def find_transitive_distance(vertices, edges):
     for i in range(inf):
         changed = False
         new_distances = {}
-        for v, d in distances.iteritems():
+        for v, d in distances.items():
             new_d = d.copy()
             new_distances[v] = new_d
-            for dest, cost in d.iteritems():
-                for leaf, cost2 in distances[dest].iteritems():
+            for dest, cost in d.items():
+                for leaf, cost2 in distances[dest].items():
                     new_cost = cost + cost2
                     old_cost = d.get(leaf, inf)
                     if new_cost < old_cost:
@@ -497,7 +499,7 @@ def get_transitive_colourer(colours, n_vertices):
         n = 1 + int(n_vertices ** 0.5)
 
         def f(link):
-            return scale[min(link * m / n, m - 1)]
+            return scale[min(link * m // n, m - 1)]
 
     else:
         def f(link):
@@ -510,7 +512,8 @@ def distance_matrix(vertices, edges,
                     utf8=False,
                     colour=None,
                     shorten_names=False,
-                    generate_key=False):
+                    generate_key=False,
+                    grouping_function=None):
     lines = []
     write = lines.append
 
@@ -527,9 +530,22 @@ def distance_matrix(vertices, edges,
 
     colours = COLOUR_SETS[colour]
 
+    colour_cycle = cycle(colours.get('alternate rows', ('',)))
+
     if vertices is None:
         vertices = sorted(set(x[0] for x in edges) | set(x[1] for x in edges))
 
+    if grouping_function is not None:
+        # we sort and colour according to the grouping function
+        # which can be used to e.g. alternate colours by site.
+        vertices = sorted(vertices, key=grouping_function)
+        colour_list = []
+        for k, v in groupby(vertices, key=grouping_function):
+            c = next(colour_cycle)
+            colour_list.extend(c for x in v)
+    else:
+        colour_list = [next(colour_cycle) for v in vertices]
+
     if shorten_names:
         edges, vertices, replacements = shorten_vertex_names(edges,
                                                              vertices,
@@ -539,7 +555,6 @@ def distance_matrix(vertices, edges,
     vlen = max(6, max(len(v) for v in vertices))
 
     # first, the key for the columns
-    colour_cycle = colours.get('alternate rows', ('',))
     c_header = colours.get('header', '')
     c_disconn = colours.get('disconnected', '')
     c_conn = colours.get('connected', '')
@@ -555,7 +570,7 @@ def distance_matrix(vertices, edges,
                                        c_reset))
     for i, v in enumerate(vertices):
         j = len(vertices) - i
-        c = colour_cycle[i % len(colour_cycle)]
+        c = colour_list[i]
         if j == 1:
             start = '%s%ssource%s' % (vspace[:-6], c_header, c_reset)
         else:
@@ -574,7 +589,7 @@ def distance_matrix(vertices, edges,
     connections = find_transitive_distance(vertices, edges)
 
     for i, v in enumerate(vertices):
-        c = colour_cycle[i % len(colour_cycle)]
+        c = colour_list[i]
         links = connections[v]
         row = []
         for v2 in vertices:
@@ -595,13 +610,14 @@ def distance_matrix(vertices, edges,
         write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
                                 ''.join(row), c_reset))
 
+    example_c = next(colour_cycle)
     if shorten_names:
         write('')
         for substitute, original in reversed(replacements):
-            write("'%s%s%s' stands for '%s%s%s'" % (colour_cycle[0],
+            write("'%s%s%s' stands for '%s%s%s'" % (example_c,
                                                     substitute,
                                                     c_reset,
-                                                    colour_cycle[0],
+                                                    example_c,
                                                     original,
                                                     c_reset))
     if generate_key:
@@ -610,7 +626,7 @@ def distance_matrix(vertices, edges,
               "indicated number of steps." % (c_header, c_reset,
                                               c_header, c_reset))
         write("%s%s%s means zero steps (it is the same DC)" %
-              (colour_cycle[0], diagonal, c_reset))
+              (example_c, diagonal, c_reset))
         write("%s1%s means a direct link" % (c_conn, c_reset))
         write("%s2%s means a transitive link involving two steps "
               "(i.e. one intermediate DC)" %