diff --git a/test/utils/test_formattree.py b/test/utils/test_formattree.py index 5e9c93c670b753782d3ebe9f8dc764c71ecaffaa..89e1fa3c0af3cd58f6a859a1c948715e4bf5b60c 100644 --- a/test/utils/test_formattree.py +++ b/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', [ diff --git a/treevalue/utils/formattree.py b/treevalue/utils/formattree.py index fc53c95f73bd1d17ee70fc6617327c14128e921d..60bfed2dfc571cbba209e1494224aa8a01c82f24 100644 --- a/treevalue/utils/formattree.py +++ b/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) + ))