lib/astree.rb in astree-1.0.0 vs lib/astree.rb in astree-1.1.0

@@ -1,8 +1,11 @@
-require 'astree/version'
 require 'colorize'
 
+Dir[File.join(__dir__, '**/*.rb')].each do |f|
+  require f
+end
+
 class ASTree
   SPACE_SIZE = 7
   T_LINE = '├─────'.freeze
   I_LINE = '│'.freeze
   L_LINE = '└─────'.freeze
@@ -20,37 +23,42 @@
     @buffer
   end
 
   private
 
-  def traverse(node, parent: nil)
-    return stringify_element(node) unless node.instance_of?(RubyVM::AbstractSyntaxTree::Node)
-
+  def traverse(node)
     buffer = stringify_node(node)
 
     children = node.children
+    children_count = children.size
 
-    until children.empty?
-      child = children.shift
-      last_element = children.empty?
-      child_buffer = traverse(child, parent: node).lines
+    children_count.times do |index|
+      child = children[index]
+      last_element = children_count == index + 1
+
+      child_buffer = if child.instance_of?(RubyVM::AbstractSyntaxTree::Node)
+                       traverse(child)
+                     else
+                       pretty_element(node: node, index: index)
+                     end.lines
+
       buffer << draw_line(token: child_buffer.shift, last_element: last_element)
 
       child_buffer.each do |line|
         buffer << draw_space(last_element: last_element) + line
       end
     end
 
     buffer
   end
 
-  def stringify_element(element)
-    str = element.inspect
-    (element.is_a?(Symbol) ? str.colorize(:magenta) : str) + "\n"
+  def pretty_element(node:, index:)
+    klass = PrettyNode.const_get(node.type)
+    klass.new(node).stringify_element(index)
   end
 
   def stringify_node(node)
-    "<%s>\n" % [node.type.to_s.colorize(:green)]
+    "<%s> [%d:%d-%d:%d]\n" % [node.type.to_s.colorize(:blue), node.first_lineno, node.first_column, node.last_lineno, node.last_column]
   end
 
   def draw_space(last_element:)
     last_element ? ' ' * SPACE_SIZE : I_LINE + ' ' * (SPACE_SIZE - 1)
   end