Document print_tree IO default (#69)

Clarify that print_tree writes to  by default when io: is not
provided, and keep the traversal helper docs consistent.

Author: evolve75

API-CHANGES.md

@@ -21,6 +21,11 @@ smooth transition to the new APIs.
 * Added [Tree::TreeNode#cmp][cmp] to compare nodes using traversal or
   relationship policies without changing the name-based `<=>` behavior.
 
+* [Tree::TreeNode#print_tree][print_tree] now accepts an `io:` keyword to
+  redirect output (defaults to `$stdout`), and
+  [Tree::TreeNode#print_tree_to_s][print_tree_to_s] returns the formatted
+  output as a string.
+
 ## Release 2.2.0 Changes
 
 * [Tree::TreeNode#add][add] now raises `ArgumentError` when attempting to add
@@ -194,6 +199,8 @@ smooth transition to the new APIs.
 [out_degree]: rdoc-ref:Tree::Utils::TreeMetricsHandler#out_degree
 [postordered_each]: rdoc-ref:Tree::TreeNode#postordered_each
 [preordered_each]: rdoc-ref:Tree::TreeNode#preordered_each
+[print_tree]: rdoc-ref:Tree::TreeNode#print_tree
+[print_tree_to_s]: rdoc-ref:Tree::TreeNode#print_tree_to_s
 [previous_sibling]: rdoc-ref:Tree::TreeNode#previous_sibling
 [remove_all]: rdoc-ref:Tree::TreeNode#remove_all!
 [rename_child]: rdoc-ref:Tree::TreeNode#rename_child

lib/tree/utils/traversal_methods.rb

@@ -32,7 +32,9 @@
 module Tree
   module Utils
     # Provides traversal helpers for TreeNode.
+    # rubocop:disable Metrics/ModuleLength
     module TreeTraversalHandler
+      require 'stringio'
       # Traverses each node (including this node) of the (sub)tree rooted at this
       # node by yielding the nodes to the specified block.
       def each # :yields: node
@@ -133,10 +135,29 @@ def each_level
       end
 
       # Pretty prints the (sub)tree rooted at this node.
+      # Output defaults to +$stdout+ unless an +io:+ is provided.
       def print_tree(level = node_depth, max_depth = nil,
-                     block = lambda { |node, prefix|
-                               puts "#{prefix} #{node.name}"
-                             })
+                     block = nil, io: $stdout, &custom_block)
+        block = resolve_print_block(block, custom_block, io)
+        block.call(self, tree_prefix(level))
+
+        # Exit if the max level is defined, and reached.
+        return unless max_depth.nil? || level < max_depth
+
+        # Child might be 'nil'
+        children do |child|
+          child&.print_tree(level + 1, max_depth, block, io: io)
+        end
+      end
+
+      # Returns the pretty-printed tree output as a string.
+      def print_tree_to_s(level = node_depth, max_depth = nil)
+        buffer = StringIO.new
+        print_tree(level, max_depth, io: buffer)
+        buffer.string
+      end
+
+      def tree_prefix(level)
         prefix = ''.dup # dup must be invoked to make this mutable.
 
         if root?
@@ -149,16 +170,19 @@ def print_tree(level = node_depth, max_depth = nil,
           prefix << (children? ? '+' : '>')
         end
 
-        block.call(self, prefix)
+        prefix
+      end
 
-        # Exit if the max level is defined, and reached.
-        return unless max_depth.nil? || level < max_depth
+      def resolve_print_block(block, custom_block, io)
+        return block if block
+        return custom_block if custom_block
 
-        # Child might be 'nil'
-        children do |child|
-          child&.print_tree(level + 1, max_depth, block)
-        end
+        lambda { |node, prefix|
+          io.puts "#{prefix} #{node.name}"
+        }
       end
+      private :tree_prefix, :resolve_print_block
     end
+    # rubocop:enable Metrics/ModuleLength
   end
 end