dev(hansbug): update the tree print system

test/utils/test_formattree.py

@@ -34,8 +34,8 @@ from treevalue.utils import (
 @pytest.mark.unittest
 class TestFormatTree(TestCase):
 
-    def format_tree(self, tree):
-        return format_tree(tree, itemgetter(0), itemgetter(1))
+    def format_tree(self, tree, encoding='utf8'):
+        return format_tree(tree, itemgetter(0), itemgetter(1), encoding)
 
     def test_single_node_tree(self):
         tree = ('foo', [])
@@ -60,6 +60,22 @@ class TestFormatTree(TestCase):
         └── qux
         '''), output)
 
+    def test_single_level_tree_with_ascii(self):
+        tree = (
+            'foo', [
+                ('bar', []),
+                ('baz', []),
+                ('qux', []),
+            ],
+        )
+        output = self.format_tree(tree, encoding='ascii')
+        self.assertEqual(dedent(u'''\
+        foo
+        +-- bar
+        +-- baz
+        `-- qux
+        '''), output)
+
     def test_multi_level_tree(self):
         tree = (
             'foo', [

treevalue/utils/formattree.py

@@ -35,43 +35,45 @@ Changes:
 
 import itertools
 import os
+import sys
 
-FORK = u'\u251c'
-LAST = u'\u2514'
-VERTICAL = u'\u2502'
-HORIZONTAL = u'\u2500'
-NEWLINE = u''
+_DEFAULT_ENCODING = os.environ.get("PYTHONIOENCODING", sys.getdefaultencoding())
 
+_UTF8_CHARS = (u'\u251c', u'\u2514', u'\u2502', u'\u2500', u'')
+_ASCII_CHARS = (u'+', u'`', u'|', u'-', u'')
 
-def _format_newlines(prefix, formatted_node):
+
+def _format_newlines(prefix, formatted_node, chars: tuple):
     """
     Convert newlines into U+23EC characters, followed by an actual newline and
     then a tree prefix so as to position the remaining text under the previous
     line.
     """
+    FORK, LAST, VERTICAL, HORIZONTAL, NEWLINE = chars
     replacement = u''.join([NEWLINE, os.linesep, prefix])
     return replacement.join(formatted_node.splitlines())
 
 
-def _format_tree(node, format_node, get_children, prefix=u''):
+def _format_tree(node, format_node, get_children, prefix=u'', chars: tuple = _UTF8_CHARS):
+    FORK, LAST, VERTICAL, HORIZONTAL, NEWLINE = chars
     children = list(get_children(node))
     next_prefix = u''.join([prefix, VERTICAL, u'   '])
     for child in children[:-1]:
         yield u''.join([
             prefix, FORK, HORIZONTAL, HORIZONTAL, u' ',
-            _format_newlines(next_prefix, format_node(child))])
-        for result in _format_tree(child, format_node, get_children, next_prefix):
+            _format_newlines(next_prefix, format_node(child), chars)])
+        for result in _format_tree(child, format_node, get_children, next_prefix, chars):
             yield result
     if children:
         last_prefix = u''.join([prefix, u'    '])
         yield u''.join([
             prefix, LAST, HORIZONTAL, HORIZONTAL, u' ',
-            _format_newlines(last_prefix, format_node(children[-1]))])
-        for result in _format_tree(children[-1], format_node, get_children, last_prefix):
+            _format_newlines(last_prefix, format_node(children[-1]), chars)])
+        for result in _format_tree(children[-1], format_node, get_children, last_prefix, chars):
             yield result
 
 
-def format_tree(node, format_node, get_children) -> str:
+def format_tree(node, format_node, get_children, encoding=None) -> str:
     r"""
     Overview:
         Format the given tree.
@@ -80,6 +82,8 @@ def format_tree(node, format_node, get_children) -> str:
         - node: Node object
         - format_node: Format node getter
         - get_children: Children getter.
+        - encoding: Encoding to be used. Default is ``None`` which means system encoding. \
+            When ASCII encoding is used, ASCII chars will be used instead of original chars.
 
     Returns:
         - formatted: Formatted string.
@@ -111,9 +115,13 @@ def format_tree(node, format_node, get_children) -> str:
             └── c
                 d
     """
-    lines = itertools.chain(
+    if 'ASCII' in (encoding or _DEFAULT_ENCODING).upper():
+        _chars = _ASCII_CHARS
+    else:
+        _chars = _UTF8_CHARS
+
+    return os.linesep.join(itertools.chain(
         [format_node(node)],
-        _format_tree(node, format_node, get_children),
+        _format_tree(node, format_node, get_children, u'', _chars),
         [u''],
-    )
-    return os.linesep.join(lines)
+    ))