Revamped pprint_tree

vallentin · andialbrecht · commit d02339582f15 · 2019-01-07T10:46:39.000+01:00
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -159,19 +159,23 @@ def __getitem__(self, item):
     def _get_repr_name(self):
         return type(self).__name__
 
-    def _pprint_tree(self, max_depth=None, depth=0, f=None):
+    def _pprint_tree(self, max_depth=None, depth=0, f=None, _pre=''):
         """Pretty-print the object tree."""
-        indent = u' | ' * depth
+        token_count = len(self.tokens)
         for idx, token in enumerate(self.tokens):
             cls = token._get_repr_name()
             value = token._get_repr_value()
 
+            last = idx == (token_count - 1)
+            pre = u'`- ' if last else u'|- '
+
             q = u'"' if value.startswith("'") and value.endswith("'") else u"'"
-            print(u"{indent}{idx:2d} {cls} {q}{value}{q}"
+            print(u"{_pre}{pre}{idx} {cls} {q}{value}{q}"
                   .format(**locals()), file=f)
 
             if token.is_group and (max_depth is None or depth < max_depth):
-                token._pprint_tree(max_depth, depth + 1, f)
+                parent_pre = u'   ' if last else u'|  '
+                token._pprint_tree(max_depth, depth + 1, f, _pre + parent_pre)
 
     def get_token_at_offset(self, offset):
         """Returns the token that is on position offset."""
diff --git a/tests/test_parse.py b/tests/test_parse.py
@@ -313,59 +313,59 @@ def test_pprint():
 
     p._pprint_tree(f=output)
     pprint = '\n'.join([
-        " 0 DML 'select'",
-        " 1 Whitespace ' '",
-        " 2 IdentifierList 'a0, b0...'",
-        " |  0 Identifier 'a0'",
-        " |  |  0 Name 'a0'",
-        " |  1 Punctuation ','",
-        " |  2 Whitespace ' '",
-        " |  3 Identifier 'b0'",
-        " |  |  0 Name 'b0'",
-        " |  4 Punctuation ','",
-        " |  5 Whitespace ' '",
-        " |  6 Identifier 'c0'",
-        " |  |  0 Name 'c0'",
-        " |  7 Punctuation ','",
-        " |  8 Whitespace ' '",
-        " |  9 Identifier 'd0'",
-        " |  |  0 Name 'd0'",
-        " | 10 Punctuation ','",
-        " | 11 Whitespace ' '",
-        " | 12 Float 'e0'",
-        " 3 Whitespace ' '",
-        " 4 Keyword 'from'",
-        " 5 Whitespace ' '",
-        " 6 Identifier '(selec...'",
-        " |  0 Parenthesis '(selec...'",
-        " |  |  0 Punctuation '('",
-        " |  |  1 DML 'select'",
-        " |  |  2 Whitespace ' '",
-        " |  |  3 Wildcard '*'",
-        " |  |  4 Whitespace ' '",
-        " |  |  5 Keyword 'from'",
-        " |  |  6 Whitespace ' '",
-        " |  |  7 Identifier 'dual'",
-        " |  |  |  0 Name 'dual'",
-        " |  |  8 Punctuation ')'",
-        " |  1 Whitespace ' '",
-        " |  2 Identifier 'q0'",
-        " |  |  0 Name 'q0'",
-        " 7 Whitespace ' '",
-        " 8 Where 'where ...'",
-        " |  0 Keyword 'where'",
-        " |  1 Whitespace ' '",
-        " |  2 Comparison '1=1'",
-        " |  |  0 Integer '1'",
-        " |  |  1 Comparison '='",
-        " |  |  2 Integer '1'",
-        " |  3 Whitespace ' '",
-        " |  4 Keyword 'and'",
-        " |  5 Whitespace ' '",
-        " |  6 Comparison '2=2'",
-        " |  |  0 Integer '2'",
-        " |  |  1 Comparison '='",
-        " |  |  2 Integer '2'",
+        "|- 0 DML 'select'",
+        "|- 1 Whitespace ' '",
+        "|- 2 IdentifierList 'a0, b0...'",
+        "|  |- 0 Identifier 'a0'",
+        "|  |  `- 0 Name 'a0'",
+        "|  |- 1 Punctuation ','",
+        "|  |- 2 Whitespace ' '",
+        "|  |- 3 Identifier 'b0'",
+        "|  |  `- 0 Name 'b0'",
+        "|  |- 4 Punctuation ','",
+        "|  |- 5 Whitespace ' '",
+        "|  |- 6 Identifier 'c0'",
+        "|  |  `- 0 Name 'c0'",
+        "|  |- 7 Punctuation ','",
+        "|  |- 8 Whitespace ' '",
+        "|  |- 9 Identifier 'd0'",
+        "|  |  `- 0 Name 'd0'",
+        "|  |- 10 Punctuation ','",
+        "|  |- 11 Whitespace ' '",
+        "|  `- 12 Float 'e0'",
+        "|- 3 Whitespace ' '",
+        "|- 4 Keyword 'from'",
+        "|- 5 Whitespace ' '",
+        "|- 6 Identifier '(selec...'",
+        "|  |- 0 Parenthesis '(selec...'",
+        "|  |  |- 0 Punctuation '('",
+        "|  |  |- 1 DML 'select'",
+        "|  |  |- 2 Whitespace ' '",
+        "|  |  |- 3 Wildcard '*'",
+        "|  |  |- 4 Whitespace ' '",
+        "|  |  |- 5 Keyword 'from'",
+        "|  |  |- 6 Whitespace ' '",
+        "|  |  |- 7 Identifier 'dual'",
+        "|  |  |  `- 0 Name 'dual'",
+        "|  |  `- 8 Punctuation ')'",
+        "|  |- 1 Whitespace ' '",
+        "|  `- 2 Identifier 'q0'",
+        "|     `- 0 Name 'q0'",
+        "|- 7 Whitespace ' '",
+        "`- 8 Where 'where ...'",
+        "   |- 0 Keyword 'where'",
+        "   |- 1 Whitespace ' '",
+        "   |- 2 Comparison '1=1'",
+        "   |  |- 0 Integer '1'",
+        "   |  |- 1 Comparison '='",
+        "   |  `- 2 Integer '1'",
+        "   |- 3 Whitespace ' '",
+        "   |- 4 Keyword 'and'",
+        "   |- 5 Whitespace ' '",
+        "   `- 6 Comparison '2=2'",
+        "      |- 0 Integer '2'",
+        "      |- 1 Comparison '='",
+        "      `- 2 Integer '2'",
         ""])
     assert output.getvalue() == pprint
 
