# vim: set fileencoding=utf-8 : from node import Node, Leaf def generate_graph(root, node_type, separator=' '): """ 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) >>> print generate_graph(n0, Node) + ╭┴╮ 0 1 >>> l2 = Leaf(2) >>> 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) node_width[node] = width return width def node_middle(node): node_len = len(node) node_mid = node_len / 2 if node_len % 2 == 1: node_mid += 1 middle = 0 for i, child in enumerate(node): if i == node_mid: break middle += separator_len + node_width[child] return middle - separator_len 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] line_width = node_width[node] lines = format_lines(node[0]) for child in node[1:]: pos = len(lines[0]) child_lines = format_lines(child) child_lines_len = len(child_lines) # A node cannot have zero child lines. assert child_lines for i, line in enumerate(lines): #print 'lines -> %d, "%s"' % (i, line) if i < child_lines_len: padding_right = ' ' * (line_width - pos \ - len(child_lines[i]) \ - separator_len) lines[i] += separator + child_lines[i] + padding_right else: # There are no more neighbor node on the right. lines[i] += ' ' * (line_width - pos) # Add the child nodes that do not have neighbor nodes on the left. for i, line in enumerate(child_lines[i+1:]): #print 'child_lines[i:] -> %d, "%s"' % (i, line) line = ' ' * (pos + separator_len) + line \ + ' ' * (line_width - separator_len - pos - len(line)) lines.append(line) # Validate that each line has an equal width try: for line in lines: assert len(line) == line_width except AssertionError: for l in lines: l = l.encode('utf-8') print l.replace(' ', '#'), len(l) l = line.encode('utf-8') print 'failed at:', l, len(l) print 'current node:', node.title(), 'width:', line_width raise # Place the title above the middle two child nodes, or when there is a # odd number of nodes, above the child node in the middle. middle = node_middle(node) #print 'node:', node.title(), 'middle:', middle, 'node_mid:', node_mid title_line = center_text(node.title(), line_width, middle) node_len = len(node) node_mid = node_len / 2 edge_line = '' for i, child in enumerate(node): if isinstance(child, node_type): middle = node_middle(child) else: middle = 0 if i < node_mid: marker = ('╭' if i == 0 else '┬') edge_line += center_text(marker.decode('utf-8'), node_width[child], middle, right='─'.decode('utf-8')) else: if i == node_mid: edge_line += '┴'.decode('utf-8') marker = ('╮' if i == node_len - 1 else '┬') edge_line += center_text(marker.decode('utf-8'), node_width[child], middle, left='─'.decode('utf-8')) #try: # assert len(title_line) == len(edge_line) #except AssertionError: # print 'title_line:', title_line, len(title_line) # print 'edge_line:', edge_line, len(edge_line) # raise # Add the line of this node before all child lines. return [title_line, edge_line] + lines calculate_width(root) lines = format_lines(root) return '\n'.join(lines).encode('utf-8') def center_text(text, width, middle=0, left=' ', right=' '): """ >>> center_text('+', 15, 11) ' + ' """ text_len = len(text) text_mid = text_len / 2 # 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)