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