python/samba/graph.py

   1 # -*- coding: utf-8 -*-
   2 # Graph topology utilities and dot file generation
   3 #
   4 # Copyright (C) Andrew Bartlett 2018.
   5 #
   6 # Written by Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
   7 #
   8 # This program is free software; you can redistribute it and/or modify
   9 # it under the terms of the GNU General Public License as published by
  10 # the Free Software Foundation; either version 3 of the License, or
  11 # (at your option) any later version.
  12 #
  13 # This program is distributed in the hope that it will be useful,
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16 # GNU General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  20
  21 from samba import colour
  22 import sys
  23 from itertools import cycle, groupby
  24
  25 FONT_SIZE = 10
  26
  27
  28 def reformat_graph_label(s):
  29     """Break DNs over multiple lines, for better shaped and arguably more
  30     readable nodes. We try to split after commas, and if necessary
  31     after hyphens or failing that in arbitrary places."""
  32     if len(s) < 12:
  33         return s
  34
  35     s = s.replace(',', ',\n')
  36     pieces = []
  37     for p in s.split('\n'):
  38         while len(p) > 20:
  39             if '-' in p[2:20]:
  40                 q, p = p.split('-', 1)
  41             else:
  42                 n = len(p) // 12
  43                 b = len(p) // n
  44                 q, p = p[:b], p[b:]
  45             pieces.append(q + '-')
  46         if p:
  47             pieces.append(p)
  48
  49     return '\\n'.join(pieces)
  50
  51
  52 def quote_graph_label(s, reformat=False):
  53     """Escape a string as graphvis requires."""
  54     # escaping inside quotes is simple in dot, because only " is escaped.
  55     # there is no need to count backslashes in sequences like \\\\"
  56     s = s.replace('"', '\"')
  57     if reformat:
  58         s = reformat_graph_label(s)
  59     return "%s" % s
  60
  61
  62 def shorten_vertex_names(vertices, suffix=',...', aggressive=False):
  63     """Replace the common suffix (in practice, the base DN) of a number of
  64     vertices with a short string (default ",..."). If this seems
  65     pointless because the replaced string is very short or the results
  66     seem strange, the original vertices are retained.
  67
  68     :param vertices: a sequence of vertices to shorten
  69     :param suffix: the replacement string [",..."]
  70     :param aggressive: replace certain common non-suffix strings
  71
  72     :return: tuple of (rename map, replacements)
  73
  74     The rename map is a dictionary mapping the old vertex names to
  75     their shortened versions. If no changes are made, replacements
  76     will be empty.
  77     """
  78     vmap = dict((v, v) for v in vertices)
  79     replacements = []
  80
  81     if len(vmap) > 1:
  82         # walk backwards along all the strings until we meet a character
  83         # that is not shared by all.
  84         i = -1
  85         vlist = list(vmap.values())
  86         try:
  87             while True:
  88                 c = set(x[i] for x in vlist)
  89                 if len(c) > 1 or '*' in c:
  90                     break
  91                 i -= 1
  92         except IndexError:
  93             # We have indexed beyond the start of a string, which should
  94             # only happen if one node is a strict suffix of all others.
  95             return vmap, replacements
  96
  97         # add one to get to the last unanimous character.
  98         i += 1
  99
 100         # now, we actually really want to split on a comma. So we walk
 101         # back to a comma.
 102         x = vlist[0]
 103         while i < len(x) and x[i] != ',':
 104             i += 1
 105
 106         if i >= -len(suffix):
 107             # there is nothing to gain here
 108             return vmap, replacements
 109
 110         replacements.append((suffix, x[i:]))
 111
 112         for k, v in vmap.items():
 113             vmap[k] = v[:i] + suffix
 114
 115     if aggressive:
 116         # Remove known common annoying strings
 117         for v in vmap.values():
 118             if ',CN=Servers,' not in v:
 119                 break
 120         else:
 121             vmap = dict((k, v.replace(',CN=Servers,', ',**,', 1))
 122                         for k, v in vmap.items())
 123             replacements.append(('**', 'CN=Servers'))
 124
 125         for v in vmap.values():
 126             if not v.startswith('CN=NTDS Settings,'):
 127                 break
 128         else:
 129             vmap = dict((k, v.replace('CN=NTDS Settings,', '*,', 1))
 130                         for k, v in vmap.items())
 131             replacements.append(('*', 'CN=NTDS Settings'))
 132
 133     return vmap, replacements
 134
 135
 136 def compile_graph_key(key_items, nodes_above=[], elisions=None,
 137                       prefix='key_', width=2):
 138     """Generate a dot file snippet that acts as a legend for a graph.
 139
 140     :param key_items: sequence of items (is_vertex, style, label)
 141     :param nodes_above: list of vertices (pushes key into right position)
 142     :param elision: tuple (short, full) indicating suffix replacement
 143     :param prefix: string used to generate key node names ["key_"]
 144     :param width: default width of node lines
 145
 146     Each item in key_items is a tuple of (is_vertex, style, label).
 147     is_vertex is a boolean indicating whether the item is a vertex
 148     (True) or edge (False). Style is a dot style string for the edge
 149     or vertex. label is the text associated with the key item.
 150     """
 151     edge_lines = []
 152     edge_names = []
 153     vertex_lines = []
 154     vertex_names = []
 155     order_lines = []
 156     for i, item in enumerate(key_items):
 157         is_vertex, style, label = item
 158         tag = '%s%d_' % (prefix, i)
 159         label = quote_graph_label(label)
 160         name = '%s_label' % tag
 161
 162         if is_vertex:
 163             order_lines.append(name)
 164             vertex_names.append(name)
 165             vertex_lines.append('%s[label="%s"; %s]' %
 166                                 (name, label, style))
 167         else:
 168             edge_names.append(name)
 169             e1 = '%se1' % tag
 170             e2 = '%se2' % tag
 171             order_lines.append(name)
 172             edge_lines.append('subgraph cluster_%s {' % tag)
 173             edge_lines.append('%s[label=src; color="#000000"; group="%s_g"]' %
 174                               (e1, tag))
 175             edge_lines.append('%s[label=dest; color="#000000"; group="%s_g"]' %
 176                               (e2, tag))
 177             edge_lines.append('%s -> %s [constraint = false; %s]' % (e1, e2,
 178                                                                      style))
 179             edge_lines.append(('%s[shape=plaintext; style=solid; width=%f; '
 180                                'label="%s\\r"]') %
 181                               (name, width, label))
 182             edge_lines.append('}')
 183
 184     elision_str = ''
 185     if elisions:
 186         for i, elision in enumerate(reversed(elisions)):
 187             order_lines.append('elision%d' % i)
 188             short, long = elision
 189             if short[0] == ',' and long[0] == ',':
 190                 short = short[1:]
 191                 long = long[1:]
 192             elision_str += ('\nelision%d[shape=plaintext; style=solid; '
 193                             'label="\“%s”  means  “%s”\\r"]\n'
 194                             % ((i, short, long)))
 195
 196     above_lines = []
 197     if order_lines:
 198         for n in nodes_above:
 199             above_lines.append('"%s" -> %s [style=invis]' %
 200                                (n, order_lines[0]))
 201
 202     s = ('subgraph cluster_key {\n'
 203          'label="Key";\n'
 204          'subgraph cluster_key_nodes {\n'
 205          'label="";\n'
 206          'color = "invis";\n'
 207          '%s\n'
 208          '}\n'
 209          'subgraph cluster_key_edges {\n'
 210          'label="";\n'
 211          'color = "invis";\n'
 212          '%s\n'
 213          '{%s}\n'
 214          '}\n'
 215          '%s\n'
 216          '}\n'
 217          '%s\n'
 218          '%s [style=invis; weight=9]'
 219          '\n'
 220          % (';\n'.join(vertex_lines),
 221             '\n'.join(edge_lines),
 222             ' '.join(edge_names),
 223             elision_str,
 224             ';\n'.join(above_lines),
 225             ' -> '.join(order_lines),
 226             ))
 227
 228     return s
 229
 230
 231 def dot_graph(vertices, edges,
 232               directed=False,
 233               title=None,
 234               reformat_labels=True,
 235               vertex_colors=None,
 236               edge_colors=None,
 237               edge_labels=None,
 238               vertex_styles=None,
 239               edge_styles=None,
 240               graph_name=None,
 241               shorten_names=False,
 242               key_items=None,
 243               vertex_clusters=None):
 244     """Generate a Graphviz representation of a list of vertices and edges.
 245
 246     :param vertices: list of vertex names (optional).
 247     :param edges:    list of (vertex, vertex) pairs
 248     :param directed: bool: whether the graph is directed
 249     :param title: optional title for the graph
 250     :param reformat_labels: whether to wrap long vertex labels
 251     :param vertex_colors: if not None, a sequence of colours for the vertices
 252     :param edge_colors: if not None, colours for the edges
 253     :param edge_labels: if not None, labels for the edges
 254     :param vertex_styles: if not None, DOT style strings for vertices
 255     :param edge_styles: if not None, DOT style strings for edges
 256     :param graph_name: if not None, name of graph
 257     :param shorten_names: if True, remove common DN suffixes
 258     :param key: (is_vertex, style, description) tuples
 259     :param vertex_clusters: list of subgraph cluster names
 260
 261     Colour, style, and label lists must be the same length as the
 262     corresponding list of edges or vertices (or None).
 263
 264     Colours can be HTML RGB strings ("#FF0000") or common names
 265     ("red"), or some other formats you don't want to think about.
 266
 267     If `vertices` is None, only the vertices mentioned in the edges
 268     are shown, and their appearance can be modified using the
 269     vertex_colors and vertex_styles arguments. Vertices appearing in
 270     the edges but not in the `vertices` list will be shown but their
 271     styles can not be modified.
 272     """
 273     out = []
 274     write = out.append
 275
 276     if vertices is None:
 277         vertices = set(x[0] for x in edges) | set(x[1] for x in edges)
 278
 279     if shorten_names:
 280         vlist = list(set(x[0] for x in edges) |
 281                      set(x[1] for x in edges) |
 282                      set(vertices))
 283         vmap, elisions = shorten_vertex_names(vlist)
 284         vertices = [vmap[x] for x in vertices]
 285         edges = [(vmap[a], vmap[b]) for a, b in edges]
 286
 287     else:
 288         elisions = None
 289
 290     if graph_name is None:
 291         graph_name = 'A_samba_tool_production'
 292
 293     if directed:
 294         graph_type = 'digraph'
 295         connector = '->'
 296     else:
 297         graph_type = 'graph'
 298         connector = '--'
 299
 300     write('/* generated by samba */')
 301     write('%s %s {' % (graph_type, graph_name))
 302     if title is not None:
 303         write('label="%s";' % (title,))
 304     write('fontsize=%s;\n' % (FONT_SIZE))
 305     write('node[fontname=Helvetica; fontsize=%s];\n' % (FONT_SIZE))
 306
 307     prev_cluster = None
 308     cluster_n = 0
 309     quoted_vertices = []
 310     for i, v in enumerate(vertices):
 311         v = quote_graph_label(v, reformat_labels)
 312         quoted_vertices.append(v)
 313         attrs = []
 314         if vertex_clusters and vertex_clusters[i]:
 315             cluster = vertex_clusters[i]
 316             if cluster != prev_cluster:
 317                 if prev_cluster is not None:
 318                     write("}")
 319                 prev_cluster = cluster
 320                 n = quote_graph_label(cluster)
 321                 if cluster:
 322                     write('subgraph cluster_%d {' % cluster_n)
 323                     cluster_n += 1
 324                     write('style = "rounded,dotted";')
 325                     write('node [style="filled"; fillcolor=white];')
 326                     write('label = "%s";' % n)
 327
 328         if vertex_styles and vertex_styles[i]:
 329             attrs.append(vertex_styles[i])
 330         if vertex_colors and vertex_colors[i]:
 331             attrs.append('color="%s"' % quote_graph_label(vertex_colors[i]))
 332         if attrs:
 333             write('"%s" [%s];' % (v, ', '.join(attrs)))
 334         else:
 335             write('"%s";' % (v,))
 336
 337     if prev_cluster:
 338         write("}")
 339
 340     for i, edge in enumerate(edges):
 341         a, b = edge
 342         if a is None:
 343             a = "Missing source value"
 344         if b is None:
 345             b = "Missing destination value"
 346
 347         a = quote_graph_label(a, reformat_labels)
 348         b = quote_graph_label(b, reformat_labels)
 349
 350         attrs = []
 351         if edge_labels:
 352             label = quote_graph_label(edge_labels[i])
 353             attrs.append('label="%s"' % label)
 354         if edge_colors:
 355             attrs.append('color="%s"' % quote_graph_label(edge_colors[i]))
 356         if edge_styles:
 357             attrs.append(edge_styles[i])  # no quoting
 358         if attrs:
 359             write('"%s" %s "%s" [%s];' % (a, connector, b, ', '.join(attrs)))
 360         else:
 361             write('"%s" %s "%s";' % (a, connector, b))
 362
 363     if key_items:
 364         key = compile_graph_key(key_items, nodes_above=quoted_vertices,
 365                                 elisions=elisions)
 366         write(key)
 367
 368     write('}\n')
 369     return '\n'.join(out)
 370
 371
 372 COLOUR_SETS = {
 373     'ansi': {
 374         'alternate rows': (colour.DARK_WHITE, colour.BLACK),
 375         'disconnected': colour.RED,
 376         'connected': colour.GREEN,
 377         'transitive': colour.DARK_YELLOW,
 378         'header': colour.UNDERLINE,
 379         'reset': colour.C_NORMAL,
 380     },
 381     'ansi-heatmap': {
 382         'alternate rows': (colour.DARK_WHITE, colour.BLACK),
 383         'disconnected': colour.REV_RED,
 384         'connected': colour.REV_GREEN,
 385         'transitive': colour.REV_DARK_YELLOW,
 386         'header': colour.UNDERLINE,
 387         'reset': colour.C_NORMAL,
 388     },
 389     'xterm-256color': {
 390         'alternate rows': (colour.xterm_256_colour(39),
 391                            colour.xterm_256_colour(45)),
 392         # 'alternate rows': (colour.xterm_256_colour(246),
 393         #                   colour.xterm_256_colour(247)),
 394         'disconnected': colour.xterm_256_colour(124, bg=True),
 395         'connected': colour.xterm_256_colour(112),
 396         'transitive': colour.xterm_256_colour(214),
 397         'transitive scale': (colour.xterm_256_colour(190),
 398                              colour.xterm_256_colour(184),
 399                              colour.xterm_256_colour(220),
 400                              colour.xterm_256_colour(214),
 401                              colour.xterm_256_colour(208),
 402                              ),
 403         'header': colour.UNDERLINE,
 404         'reset': colour.C_NORMAL,
 405     },
 406     'xterm-256color-heatmap': {
 407         'alternate rows': (colour.xterm_256_colour(171),
 408                            colour.xterm_256_colour(207)),
 409         # 'alternate rows': (colour.xterm_256_colour(246),
 410         #                    colour.xterm_256_colour(247)),
 411         'disconnected': colour.xterm_256_colour(124, bg=True),
 412         'connected': colour.xterm_256_colour(112, bg=True),
 413         'transitive': colour.xterm_256_colour(214, bg=True),
 414         'transitive scale': (colour.xterm_256_colour(190, bg=True),
 415                              colour.xterm_256_colour(184, bg=True),
 416                              colour.xterm_256_colour(220, bg=True),
 417                              colour.xterm_256_colour(214, bg=True),
 418                              colour.xterm_256_colour(208, bg=True),
 419                              ),
 420         'header': colour.UNDERLINE,
 421         'reset': colour.C_NORMAL,
 422     },
 423     None: {
 424         'alternate rows': ('',),
 425         'disconnected': '',
 426         'connected': '',
 427         'transitive': '',
 428         'header': '',
 429         'reset': '',
 430     }
 431 }
 432
 433 CHARSETS = {
 434     'utf8': {
 435         'vertical': '│',
 436         'horizontal': '─',
 437         'corner': '╭',
 438         # 'diagonal': '╲',
 439         'diagonal': '·',
 440         # 'missing': '🕱',
 441         'missing': '-',
 442         'right_arrow': '←',
 443     },
 444     'ascii': {
 445         'vertical': '|',
 446         'horizontal': '-',
 447         'corner': ',',
 448         'diagonal': '0',
 449         'missing': '-',
 450         'right_arrow': '<-',
 451     }
 452 }
 453
 454
 455 def find_transitive_distance(vertices, edges):
 456     all_vertices = (set(vertices) |
 457                     set(e[0] for e in edges) |
 458                     set(e[1] for e in edges))
 459
 460     if all_vertices != set(vertices):
 461         print("there are unknown vertices: %s" %
 462               (all_vertices - set(vertices)),
 463               file=sys.stderr)
 464
 465     # with n vertices, we are always less than n hops away from
 466     # anywhere else.
 467     inf = len(all_vertices)
 468     distances = {}
 469     for v in all_vertices:
 470         distances[v] = {v: 0}
 471
 472     for src, dest in edges:
 473         distances[src][dest] = distances[src].get(dest, 1)
 474
 475     # This algorithm (and implementation) seems very suboptimal.
 476     # potentially O(n^4), though n is smallish.
 477     for i in range(inf):
 478         changed = False
 479         new_distances = {}
 480         for v, d in distances.items():
 481             new_d = d.copy()
 482             new_distances[v] = new_d
 483             for dest, cost in d.items():
 484                 for leaf, cost2 in distances[dest].items():
 485                     new_cost = cost + cost2
 486                     old_cost = d.get(leaf, inf)
 487                     if new_cost < old_cost:
 488                         new_d[leaf] = new_cost
 489                         changed = True
 490
 491         distances = new_distances
 492         if not changed:
 493             break
 494
 495     # filter out unwanted vertices and infinite links
 496     answer = {}
 497     for v in vertices:
 498         answer[v] = {}
 499         for v2 in vertices:
 500             a = distances[v].get(v2, inf)
 501             if a < inf:
 502                 answer[v][v2] = a
 503
 504     return answer
 505
 506
 507 def get_transitive_colourer(colours, n_vertices):
 508     if 'transitive scale' in colours:
 509         scale = colours['transitive scale']
 510         m = len(scale)
 511         n = 1 + int(n_vertices ** 0.5)
 512
 513         def f(link):
 514             if not isinstance(link, int):
 515                 return ''
 516             return scale[min(link * m // n, m - 1)]
 517
 518     else:
 519         def f(link):
 520             return colours['transitive']
 521
 522     return f
 523
 524
 525 def distance_matrix(vertices, edges,
 526                     utf8=False,
 527                     colour=None,
 528                     shorten_names=False,
 529                     generate_key=False,
 530                     grouping_function=None,
 531                     row_comments=None):
 532     lines = []
 533     write = lines.append
 534
 535     charset = CHARSETS['utf8' if utf8 else 'ascii']
 536     vertical = charset['vertical']
 537     horizontal = charset['horizontal']
 538     corner = charset['corner']
 539     diagonal = charset['diagonal']
 540     missing = charset['missing']
 541     right_arrow = charset['right_arrow']
 542
 543     colours = COLOUR_SETS[colour]
 544
 545     colour_cycle = cycle(colours.get('alternate rows', ('',)))
 546
 547     if vertices is None:
 548         vertices = sorted(set(x[0] for x in edges) | set(x[1] for x in edges))
 549
 550     if grouping_function is not None:
 551         # we sort and colour according to the grouping function
 552         # which can be used to e.g. alternate colours by site.
 553         vertices = sorted(vertices, key=grouping_function)
 554         colour_list = []
 555         for k, v in groupby(vertices, key=grouping_function):
 556             c = next(colour_cycle)
 557             colour_list.extend(c for x in v)
 558     else:
 559         colour_list = [next(colour_cycle) for v in vertices]
 560
 561     if shorten_names:
 562         vlist = list(set(x[0] for x in edges) |
 563                      set(x[1] for x in edges) |
 564                      set(vertices))
 565         vmap, replacements = shorten_vertex_names(vlist, '+',
 566                                                   aggressive=True)
 567         vertices = [vmap[x] for x in vertices]
 568         edges = [(vmap[a], vmap[b]) for a, b in edges]
 569
 570     vlen = max(6, max(len(v) for v in vertices))
 571
 572     # first, the key for the columns
 573     c_header = colours.get('header', '')
 574     c_disconn = colours.get('disconnected', '')
 575     c_conn = colours.get('connected', '')
 576     c_reset = colours.get('reset', '')
 577
 578     colour_transitive = get_transitive_colourer(colours, len(vertices))
 579
 580     vspace = ' ' * vlen
 581     verticals = ''
 582     write("%*s %s  %sdestination%s" % (vlen, '',
 583                                        ' ' * len(vertices),
 584                                        c_header,
 585                                        c_reset))
 586     for i, v in enumerate(vertices):
 587         j = len(vertices) - i
 588         c = colour_list[i]
 589         if j == 1:
 590             start = '%s%ssource%s' % (vspace[:-6], c_header, c_reset)
 591         else:
 592             start = vspace
 593         write('%s %s%s%s%s%s %s%s' % (start,
 594                                       verticals,
 595                                       c_reset,
 596                                       c,
 597                                       corner,
 598                                       horizontal * j,
 599                                       v,
 600                                       c_reset
 601                                       ))
 602         verticals += c + vertical
 603
 604     connections = find_transitive_distance(vertices, edges)
 605
 606     for i, v in enumerate(vertices):
 607         c = colour_list[i]
 608         links = connections[v]
 609         row = []
 610         for v2 in vertices:
 611             link = links.get(v2)
 612             if link is None:
 613                 row.append('%s%s' % (c_disconn, missing))
 614                 continue
 615             if link == 0:
 616                 row.append('%s%s%s%s' % (c_reset, c, diagonal, c_reset))
 617             elif link == 1:
 618                 row.append('%s1%s' % (c_conn, c_reset))
 619             else:
 620                 ct = colour_transitive(link)
 621                 if link > 9:
 622                     link = '>'
 623                 row.append('%s%s%s' % (ct, link, c_reset))
 624
 625         if row_comments is not None and row_comments[i]:
 626             row.append('%s %s %s' % (c_reset, right_arrow, row_comments[i]))
 627
 628         write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
 629                                 ''.join(row), c_reset))
 630
 631     example_c = next(colour_cycle)
 632     if shorten_names:
 633         write('')
 634         for substitute, original in reversed(replacements):
 635             write("'%s%s%s' stands for '%s%s%s'" % (example_c,
 636                                                     substitute,
 637                                                     c_reset,
 638                                                     example_c,
 639                                                     original,
 640                                                     c_reset))
 641     if generate_key:
 642         write('')
 643         write("Data can get from %ssource%s to %sdestination%s in the "
 644               "indicated number of steps." % (c_header, c_reset,
 645                                               c_header, c_reset))
 646         write("%s%s%s means zero steps (it is the same DC)" %
 647               (example_c, diagonal, c_reset))
 648         write("%s1%s means a direct link" % (c_conn, c_reset))
 649         write("%s2%s means a transitive link involving two steps "
 650               "(i.e. one intermediate DC)" %
 651               (colour_transitive(2), c_reset))
 652         write("%s%s%s means there is no connection, even through other DCs" %
 653               (c_disconn, missing, c_reset))
 654
 655     return '\n'.join(lines)
 656
 657
 658 def pad_char(char, digits, padding=' '):
 659     if digits == 1:
 660         padding = ''
 661     return ' ' * (digits - 1) + char + padding
 662
 663
 664 def transpose_dict_matrix(m):
 665     m2 = {}
 666     for k1, row in m.items():
 667         for k2, dist in row.items():
 668             m2.setdefault(k2, {})[k1] = dist
 669     return m2
 670
 671
 672 def full_matrix(rows,
 673                 utf8=False,
 674                 colour=None,
 675                 shorten_names=False,
 676                 generate_key=False,
 677                 grouping_function=None,
 678                 row_comments=None,
 679                 colour_scale=None,
 680                 digits=1,
 681                 ylabel='source',
 682                 xlabel='destination',
 683                 transpose=True):
 684     lines = []
 685     write = lines.append
 686
 687     if transpose:
 688         rows = transpose_dict_matrix(rows)
 689
 690     use_padding = digits > 1
 691
 692     charset = CHARSETS['utf8' if utf8 else 'ascii']
 693     vertical = pad_char(charset['vertical'], digits)
 694     horizontal = charset['horizontal'] * (digits + use_padding)
 695     corner = pad_char(charset['corner'], digits,
 696                       charset['horizontal'])
 697     diagonal = pad_char(charset['diagonal'], digits)
 698     missing = pad_char(charset['missing'], digits)
 699     toobig = pad_char('>', digits)
 700     right_arrow = charset['right_arrow']
 701     empty = pad_char(' ', digits)
 702
 703     colours = COLOUR_SETS[colour]
 704
 705     colour_cycle = cycle(colours.get('alternate rows', ('',)))
 706     vertices = list(rows.keys())
 707     if grouping_function is not None:
 708         # we sort and colour according to the grouping function
 709         # which can be used to e.g. alternate colours by site.
 710         vertices.sort(key=grouping_function)
 711         colour_list = []
 712         for k, v in groupby(vertices, key=grouping_function):
 713             c = next(colour_cycle)
 714             colour_list.extend(c for x in v)
 715     else:
 716         colour_list = [next(colour_cycle) for v in vertices]
 717
 718     if shorten_names:
 719         vmap, replacements = shorten_vertex_names(vertices, '+',
 720                                                   aggressive=True)
 721         rows2 = {}
 722         for vert, r in rows.items():
 723             rows2[vmap[vert]] = dict((vmap[k], v) for k, v in r.items())
 724
 725         rows = rows2
 726         vertices = list(rows.keys())
 727
 728     vlen = max(6, len(xlabel), max(len(v) for v in vertices))
 729
 730     # first, the key for the columns
 731     c_header = colours.get('header', '')
 732     c_disconn = colours.get('disconnected', '')
 733     c_conn = colours.get('connected', '')
 734     c_reset = colours.get('reset', '')
 735
 736     if colour_scale is None:
 737         colour_scale = len(rows)
 738     colour_transitive = get_transitive_colourer(colours, colour_scale)
 739
 740     vspace = ' ' * vlen
 741     verticals = ''
 742     write("%s %s %s%s%s" % (vspace,
 743                             empty * (len(rows) + 1),
 744                             c_header,
 745                             xlabel,
 746                             c_reset))
 747     for i, v in enumerate(vertices):
 748         j = len(rows) - i
 749         c = colour_list[i]
 750         if j == 1:
 751             start = '%s%s%s%s' % (vspace[:-len(ylabel)],
 752                                   c_header,
 753                                   ylabel,
 754                                   c_reset)
 755         else:
 756             start = vspace
 757         write('%s %s%s%s%s%s %s%s' % (start,
 758                                       verticals,
 759                                       c_reset,
 760                                       c,
 761                                       corner,
 762                                       horizontal * j,
 763                                       v,
 764                                       c_reset
 765                                       ))
 766         verticals += '%s%s' % (c, vertical)
 767
 768     end_cell = '%s%s' % (' ' * use_padding, c_reset)
 769     overflow = False
 770     for i, v in enumerate(vertices):
 771         links = rows[v]
 772         c = colour_list[i]
 773         row = []
 774         for v2 in vertices:
 775             if v2 not in links:
 776                 row.append('%s%s%s' % (c_disconn, missing, c_reset))
 777             elif v == v2:
 778                 row.append('%s%s%s%s' % (c_reset, c, diagonal, c_reset))
 779             else:
 780                 link = links[v2]
 781                 if link >= 10 ** digits:
 782                     ct = colour_transitive(link)
 783                     row.append('%s%s%s' % (ct, toobig, c_reset))
 784                     overflow = True
 785                     continue
 786                 if link == 0:
 787                     ct = c_conn
 788                 else:
 789                     ct = colour_transitive(link)
 790                 row.append('%s%*s%s' % (ct, digits, link, end_cell))
 791
 792         if row_comments is not None and row_comments[i]:
 793             row.append('%s %s %s' % (c_reset, right_arrow, row_comments[i]))
 794
 795         write('%s%*s%s %s%s' % (c, vlen, v, c_reset,
 796                                 ''.join(row), c_reset))
 797
 798     if overflow or shorten_names:
 799         write('')
 800
 801     if overflow:
 802             write("'%s%s%s' means greater than %d " %
 803                   (colour_transitive(10 ** digits),
 804                    toobig,
 805                    c_reset,
 806                    10 ** digits - 1))
 807
 808     if shorten_names:
 809         example_c = next(colour_cycle)
 810         for substitute, original in reversed(replacements):
 811             write("'%s%s%s' stands for '%s%s%s'" % (example_c,
 812                                                     substitute,
 813                                                     c_reset,
 814                                                     example_c,
 815                                                     original,
 816                                                     c_reset))
 817
 818     return '\n'.join(lines)