Added method to draw tree structure.

doorstop/cli/main.py

@@ -259,7 +259,9 @@ def _run(args, cwd, err):  # pylint: disable=W0613
         return False
 
     if tree and valid:
-        print("valid tree: {}".format(tree))
+        print()
+        print(tree.draw())
+        print()
     return valid
 
 

doorstop/core/test/test_tree.py

@@ -26,26 +26,30 @@ class TestTreeStrings(unittest.TestCase):  # pylint: disable=R0904
     def setUpClass(cls):
         a = Tree('a', root='.')
         b1 = Tree('b1', parent=a, root='.')
+        d = Tree('d', parent=b1, root='.')
+        e = Tree('e', parent=d, root='.')
         b2 = Tree('b2', parent=a, root='.')
         c1 = Tree('c1', parent=b2, root='.')
         c2 = Tree('c2', parent=b2, root='.')
         a.children = [b1, b2]
+        b1.children = [d]
+        d.children = [e]
         b2.children = [c1, c2]
         cls.tree = a
 
     def test_repr(self):
         """Verify trees can be represented."""
-        text = "<Tree a <- [ b1, b2 <- [ c1, c2 ] ]>"
+        text = "<Tree a <- [ b1 <- [ d <- [ e ] ], b2 <- [ c1, c2 ] ]>"
         self.assertEqual(text, repr(self.tree))
 
     def test_str(self):
         """Verify trees can be converted to strings."""
-        text = "a <- [ b1, b2 <- [ c1, c2 ] ]"
+        text = "a <- [ b1 <- [ d <- [ e ] ], b2 <- [ c1, c2 ] ]"
         self.assertEqual(text, str(self.tree))
 
     def test_len(self):
         """Verify a tree lengths are correct."""
-        self.assertEqual(5, len(self.tree))
+        self.assertEqual(7, len(self.tree))
 
     def test_getitem(self):
         """Verify item access is not allowed on trees."""
@@ -54,13 +58,33 @@ def test_getitem(self):
     def test_iter(self):
         """Verify a tree can be iterated over."""
         items = [d for d in self.tree]
-        self.assertListEqual(['a', 'b1', 'b2', 'c1', 'c2'], items)
+        self.assertListEqual(['a', 'b1', 'd', 'e', 'b2', 'c1', 'c2'], items)
 
     def test_contains(self):
         """Verify a tree can be checked for contents."""
         child = self.tree.children[1].children[0]
         self.assertIn(child.document, self.tree)
 
+    def test_draw(self):
+        """Verify trees structure can be drawn."""
+        text = ("a" + '\n'
+                "|   " + '\n'
+                "├ ‒ b1" + '\n'
+                "|   |   " + '\n'
+                "|   └ ‒ d" + '\n'
+                "|       |   " + '\n'
+                "|       └ ‒ e" + '\n'
+                "|   " + '\n'
+                "└ ‒ b2" + '\n'
+                "    |   " + '\n'
+                "    ├ ‒ c1" + '\n'
+                "    |   " + '\n'
+                "    └ ‒ c2")
+        logging.debug('\n' + text)
+        text2 = self.tree.draw()
+        logging.debug('\n' + text2)
+        self.assertEqual(text, text2)
+
     @patch('doorstop.settings.REORDER', False)
     def test_from_list(self):
         """Verify a tree can be created from a list."""

doorstop/core/tree.py

@@ -80,18 +80,10 @@ def __init__(self, document, parent=None, root=None):
         self._document_cache = {}
 
     def __repr__(self):
-        return "<Tree {}>".format(self)
+        return "<Tree {}>".format(self._draw_line())
 
     def __str__(self):
-        # Build parent prefix string (enables mock testing)
-        prefix = getattr(self.document, 'prefix', self.document)
-        # Build children prefix strings
-        children = ", ".join(str(c) for c in self.children)
-        # Format the tree
-        if children:
-            return "{} <- [ {} ]".format(prefix, children)
-        else:
-            return "{}".format(prefix)
+        return self._draw_line()
 
     def __len__(self):
         if self.document:
@@ -431,6 +423,41 @@ def load(self, reload=False):
         # Set meta attributes
         self._loaded = True
 
+    def draw(self):
+        return '\n'.join(self._draw_lines())
+
+    def _draw_line(self):
+        """Get the tree structure in one line."""
+        # Build parent prefix string (`getattr` to enable mock testing)
+        prefix = getattr(self.document, 'prefix', '') or str(self.document)
+        # Build children prefix strings
+        children = ", ".join(c._draw_line() for c in self.children)  # pylint: disable=W0212
+        # Format the tree
+        if children:
+            return "{} <- [ {} ]".format(prefix, children)
+        else:
+            return "{}".format(prefix)
+
+    def _draw_lines(self):
+        """Generate lines of the tree structure."""
+        # Build parent prefix string (`getattr` to enable mock testing)
+        prefix = getattr(self.document, 'prefix', '') or str(self.document)
+        yield prefix
+        # Build child prefix strings
+        for count, child in enumerate(self.children, start=1):
+            yield '|   '
+            if count < len(self.children):
+                base = '|   '
+                indent = '├ ‒ '
+            else:
+                base = '    '
+                indent = '└ ‒ '
+            for count, line in enumerate(child._draw_lines(), start=1):  # pylint: disable=W0212
+                if count == 1:
+                    yield indent + line
+                else:
+                    yield base + line
+
     def delete(self):
         """Delete the tree and its documents and items."""
         for document in self: