Fix #139: replace stdout with end in show.

tests/test_tree.py

@@ -2,14 +2,16 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 
+import contextlib
+
 import sys
 
 import os
 
 try:
-    from StringIO import StringIO as BytesIO
+    from StringIO import StringIO
 except ImportError:
-    from io import BytesIO
+    from io import StringIO
 import unittest
 from treelib import Tree, Node
 from treelib.tree import NodeIDAbsentError, LoopError
@@ -242,31 +244,38 @@ def test_paste_tree(self):
         new_tree.create_node("Mark", "mark", parent="jill")
         self.tree.paste("jane", new_tree)
         self.assertEqual("jill" in self.tree.is_branch("jane"), True)
-        self.tree.show()
+        f = StringIO()
+        with contextlib.redirect_stdout(f):
+            self.tree.show()
         self.assertEqual(
-            self.tree._reader,
-            u'''Hárry
+            f.getvalue(),
+u'''Hárry
 ├── Bill
 │   └── George
 └── Jane
     ├── Diane
     └── Jill
         └── Mark
+
 '''
         )
+        f.close()
         self.tree.remove_node("jill")
         self.assertNotIn('jill', self.tree.nodes.keys())
         self.assertNotIn('mark', self.tree.nodes.keys())
-        self.tree.show()
+        f = StringIO()
+        with contextlib.redirect_stdout(f):
+            self.tree.show(end='\n')
         self.assertEqual(
-            self.tree._reader,
-            u'''Hárry
+            f.getvalue(),
+u'''Hárry
 ├── Bill
 │   └── George
 └── Jane
     └── Diane
 '''
         )
+        f.close()
 
     def test_paste_at_root_level(self):
         t1 = Tree()
@@ -285,15 +294,22 @@ def test_paste_at_root_level(self):
         self.assertEqual(t1.root, 'r')
         self.assertNotIn('r2', t1._nodes.keys())
         self.assertEqual(set(t1._nodes.keys()), {'r', 'a', 'a1', 'b', 'c', 'd', 'd1'})
-        t1.show()
-        self.assertEqual(t1._reader, '''root
+        f = StringIO()
+        with contextlib.redirect_stdout(f):
+            t1.show()
+        self.assertEqual(
+            f.getvalue(),
+'''root
 ├── A
 │   └── A1
 ├── B
 ├── C
 └── D
     └── D1
-''')
+
+'''
+        )
+        f.close()
 
     def test_rsearch(self):
         for nid in ["hárry", "jane", "diane"]:
@@ -340,19 +356,15 @@ def __init__(self, color):
     def test_show_data_property(self):
         new_tree = Tree()
 
-        sys.stdout = open(os.devnull, "w")  # stops from printing to console
-
-        try:
-            new_tree.show()
+        with open(os.devnull, 'w') as f:
+            with contextlib.redirect_stdout(f):
+                new_tree.show()
 
-            class Flower(object):
-                def __init__(self, color):
-                    self.color = color
-            new_tree.create_node("Jill", "jill", data=Flower("white"))
-            new_tree.show(data_property="color")
-        finally:
-            sys.stdout.close()
-            sys.stdout = sys.__stdout__  # stops from printing to console
+                class Flower(object):
+                    def __init__(self, color):
+                        self.color = color
+                new_tree.create_node("Jill", "jill", data=Flower("white"))
+                new_tree.show(data_property="color")
 
     def test_level(self):
         self.assertEqual(self.tree.level('hárry'),  0)
@@ -382,13 +394,9 @@ def test_show(self):
         if sys.version_info[0] < 3:
             reload(sys)
             sys.setdefaultencoding('utf-8')
-        sys.stdout = open(os.devnull, "w")  # stops from printing to console
-
-        try:
-            self.tree.show()
-        finally:
-            sys.stdout.close()
-            sys.stdout = sys.__stdout__  # stops from printing to console
+        with open(os.devnull, "w") as f:
+            with contextlib.redirect_stdout(f):
+                self.tree.show()
 
     def tearDown(self):
         self.tree = None

treelib/tree.py

@@ -780,7 +780,8 @@ def _write_line(line, f):
                              key, reverse, line_type, data_property, func=handler)
 
     def show(self, nid=None, level=ROOT, idhidden=True, filter=None,
-             key=None, reverse=False, line_type='ascii-ex', data_property=None, stdout=True):
+             key=None, reverse=False, line_type='ascii-ex', data_property=None,
+             end='\n\n'):
         """
         Print the tree structure in hierarchy style.
 
@@ -801,23 +802,23 @@ def show(self, nid=None, level=ROOT, idhidden=True, filter=None,
         :param reverse: the ``reverse`` param for sorting :class:`Node` objects in the same level.
         :param line_type:
         :param data_property: the property on the node data object to be printed.
+        :param end: the string to append to the final line of the output. By
+            default, an extra newline is printed; pass '\n' if this is not
+            desired.
         :return: None
         """
-        self._reader = ""
+        self._reader = []
 
         def write(line):
-            self._reader += line.decode('utf-8') + "\n"
+            self._reader.append(line.decode('utf-8'))
 
         try:
             self.__print_backend(nid, level, idhidden, filter,
                                  key, reverse, line_type, data_property, func=write)
         except NodeIDAbsentError:
             print('Tree is empty')
 
-        if stdout:
-            print(self._reader)
-        else:
-            return self._reader
+        print('\n'.join(self._reader), end=end)
 
     def siblings(self, nid):
         """