Allow Node#traverse to return enumerator. Fixes #117.

ioquatix · ioquatix · commit 85d65591d9fd · 2022-10-31T17:50:47.000+13:00
diff --git a/lib/async/node.rb b/lib/async/node.rb
@@ -209,13 +209,19 @@ def consume
 			end
 		end
 		
+		def traverse(&block)
+			return enum_for(:traverse) unless block_given?
+			
+			self.traverse_recurse(&block)
+		end
+		
 		# Traverse the tree.
 		# @yields {|node, level| ...} The node and the level relative to the given root.
-		def traverse(level = 0, &block)
+		protected def traverse_recurse(level = 0, &block)
 			yield self, level
 			
 			@children&.each do |child|
-				child.traverse(level + 1, &block)
+				child.traverse_recurse(level + 1, &block)
 			end
 		end
 		
diff --git a/test/async/node.rb b/test/async/node.rb
@@ -301,4 +301,14 @@
 			node.terminate
 		end
 	end
+	
+	with '#traverse' do
+		it 'can traverse without a block' do
+			middle = Async::Node.new(node, annotation: "middle")
+			child1 = Async::Node.new(middle, transient: true, annotation: "child1")
+			child2 = Async::Node.new(middle, annotation: "child2")
+			
+			expect(node.traverse.to_a).to be == [[node, 0], [middle, 1], [child1, 2], [child2, 2]]
+		end
+	end
 end
