index 6cdd2ef..7d195e3 100644 (file)
# 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:
@@ -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 for x in edges) | set(x 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_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]
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,
+            write("'%s%s%s' stands for '%s%s%s'" % (example_c,
substitute,
c_reset,
-                                                    colour_cycle,
+                                                    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,