# vim: set fileencoding=utf-8 : # XXX Used in doctests (we should use them in the __main__ section below too). 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') def generate_graph(root, 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. >>> from node import Leaf, Node >>> l0, l1 = Leaf(0), Leaf(1) >>> n0 = Node('+', l0, l1) >>> l2 = Leaf(2) >>> print generate_graph(n0) + ╭┴╮ 0 1 >>> n1 = Node('-', l2) >>> print generate_graph(n1) - │ 2 >>> n2 = Node('*', n0, n1) >>> print generate_graph(n2) * ╭─┴╮ + - ╭┴╮ │ 0 1 2 """ node_width = {} node_middle = {} separator_len = len(separator) def calculate_node_sizes(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 node.nodes: node_width[node] = title_len node_middle[node] = int((title_len - 1) / 2) return title_len node_len = len(node) width = 0 middle = 0 middle_pos = int(node_len / 2) for i, child in enumerate(node): tmp = calculate_node_sizes(child) if i < middle_pos: middle += tmp width += tmp middle += max(middle_pos - int(node_len % 2 == 0), 0) * separator_len # Add a separator between each node (thus n - 1 separators). width += separator_len * (node_len - 1) # If the title of the node is wider than the sum of its children, the # title's width should be used. if title_len > width: width = title_len middle = int(title_len / 2) # print 'width of "%s":' % node.title(), width node_width[node] = width node_middle[node] = middle return width def format_lines(node): if not node.nodes: # 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] # TODO: substitute box_widths with node_width box_widths = [len(lines[0]) for lines in child_lines] node_len = len(node) middle_node = int(node_len / 2) middle = node_middle[node] title_line = center_text(node.title(), line_width, middle) 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=node_middle[child], 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=node_middle[child], 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_node_sizes(root) if verbose: # pragma: nocover print '------- node_{width,middle} ---------' for node, width in node_width.iteritems(): print node.title(), 'width:', width, 'middle:', node_middle[node] 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) ' + ' >>> sep = '─'.decode('utf-8') >>> left = center_text('╭'.decode('utf-8'), 11, 8, right=sep) >>> len(left) == 11 True >>> right = center_text('╮'.decode('utf-8'), 3, 2, left=sep) >>> len(right) == 3 True >>> edge_line = left + '┴'.decode('utf-8') + right >>> len(edge_line) == 15 True >>> title_line = center_text('+', 15, 11) >>> print '|%s|\\n|%s|' % (title_line, edge_line.encode('utf-8')) | + | | ╭──┴──╮| """ 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)