# vim: set fileencoding=utf-8 : def generate_graph(root, node_type, separator=' ', verbose=False): """ Return a text-based, utf-8 graph of a tree-like structure. Each tree node is represented by a length-2 list. If a node has an attribute called ``title``, that attribute will be called. That way, the node can return a specific title, otherwise ``+`` is used. >>> l0, l1 = Leaf(0), Leaf(1) >>> n0 = Node('+', l0, l1) >>> l2 = Leaf(2) >>> print generate_graph(n0, Node) + ╭┴╮ 0 1 >>> n1 = Node('-', l2) >>> print generate_graph(n1, Node) - │ 2 >>> n2 = Node('*', n0, n1) >>> print generate_graph(n2, Node) * ╭─┴╮ + - ╭┴╮ │ 0 1 2 """ node_width = {} separator_len = len(separator) def calculate_width(node): title = node.title() title_len = len(title) # Leaves do not have children and therefore the length of its title is # the width of the leaf. if not isinstance(node, node_type): node_width[node] = title_len return title_len width = 0 for child in node: width += calculate_width(child) # Add a separator between each node (thus n - 1 separators). width += separator_len * (len(node) - 1) # Odd numbers of children should be minus 1, since the middle child # can be placed directly below the parent. With even numbers, there # is no middle child, so the space below the parent cannot be used. #if len(node) % 2 == 1: # width -= 1 # If the title of the node is wider than the sum of its children, the # title's width should be used. width = max(title_len, width) # print 'width of "%s":' % node.title(), width node_width[node] = width return width def format_lines(node): if not isinstance(node, node_type): # Leaf titles do not need to be centered, since the parent will # center those lines. And if there are no parents, the entire graph # consists of a single leaf, so in that case there still is no # reason to center it. return [node.title()] # At least one child, otherwise it would be a leaf. assert node[0] child_lines = [format_lines(child) for child in node] max_height = max(map(len, child_lines)) # Assert that all child boxes are of equal height for lines in child_lines: additional_line = separator * len(lines[0]) lines += [additional_line for i in range(max_height - len(lines))] assert len(child_lines[0]) == max_height from copy import deepcopy result = deepcopy(child_lines[0]) for lines in child_lines[1:]: assert len(lines) == max_height for i, line in enumerate(lines): result[i] += separator + line line_width = node_width[node] box_widths = [len(lines[0]) for lines in child_lines] node_len = len(node) middle_node = int(node_len / 2) middle = sum([box_widths[i] for i in range(middle_node)]) \ + max(middle_node - int(node_len % 2 == 0), 0) * separator_len title_line = center_text(node.title(), line_width, middle) pipe_sign = '│'.decode('utf-8') dash_sign = '─'.decode('utf-8') cross_sign = '┼'.decode('utf-8') tsplit_dn_sign = '┬'.decode('utf-8') tsplit_up_sign = '┴'.decode('utf-8') left_sign = '╭'.decode('utf-8') right_sign = '╮'.decode('utf-8') if node_len == 1: # Unary operators edge_line = center_text(pipe_sign, box_widths[0], middle) elif node_len % 2: # n-ary operators (n is odd) edge_line = '' for i, child in enumerate(node): if i > 0: edge_line += dash_sign if i < middle_node: marker = (left_sign if i == 0 else tsplit_dn_sign) edge_line += center_text(marker, box_widths[i], middle=0, right=dash_sign) else: if i == middle_node: marker = cross_sign edge_line += center_text(marker, box_widths[i], middle=0, right=dash_sign, left=dash_sign) else: if i == node_len - 1: marker = right_sign else: marker = tsplit_dn_sign edge_line += center_text(marker, box_widths[i], middle=0, left=dash_sign) else: # n-ary operators (n is even) edge_line = '' for i, child in enumerate(node): if i > 0: if i == middle_node: edge_line += tsplit_up_sign else: edge_line += dash_sign if i < middle_node: marker = (left_sign if i == 0 else tsplit_dn_sign) edge_line += center_text(marker, box_widths[i], middle=0, right=dash_sign) else: if i == node_len - 1: marker = right_sign else: marker = tsplit_dn_sign edge_line += center_text(marker, box_widths[i], middle=0, left=dash_sign) try: assert len(title_line) == len(edge_line) except AssertionError: # pragma: nocover print '------------------' print 'line_width:', line_width print 'title_line:', title_line, 'len:', len(title_line) print 'edge_line: %s (%d)' % (edge_line.encode('utf-8'), len(edge_line)) print 'lines:' print '\n'.join(map(lambda x: x + ' %d' % len(x), lines)) raise Exception() # Add the line of this node before all child lines. return [title_line, edge_line] + result calculate_width(root) #if verbose: # print '------- node_width ---------' # for node, width in node_width.iteritems(): # print node.title(), 'width:', width lines = format_lines(root) # Strip trailing whitespace. return '\n'.join(map(lambda x: x.rstrip(), lines)).encode('utf-8') def center_text(text, width, middle=0, left=' ', right=' '): """ >>> print center_text('│', 1, 1) │ >>> center_text('+', 15, 11) ' + ' """ text_len = len(text) text_mid = text_len / 2 #print '---------' #print 'text_len:', text_len #print 'text_mid:', text_mid #print 'width:', width #print 'middle:', middle #print '---------' # TODO: this code requires cleanup. if middle: # If this is true, the text is at the left. if text_mid > middle: text += left * (width - text_len) # If this is true, the text is at the right. elif text_mid > (width - middle): text = right * (width - text_len) + text # Else, the text has spacing on its left and right. else: text = left * (middle - text_mid) + text text += right * (width - len(text)) return text spacing = width - text_len # Even number of spaces can be split equally. if spacing % 2 == 0: return left * (spacing / 2) + text + right * (spacing / 2) # For an odd number of space, put the extra space at the end. return left * (spacing / 2) + text + right * (spacing / 2 + 1)