RubyVM::AbstractSyntaxTree で Node を辿る処理をシンプルに実装する

Ruby の Enumerator を使うと簡単にRubyのコードを処理できて便利だったので、ブログに書いておく。

ast = RubyVM::AbstractSyntaxTree.parse(<<~RUBY)
  class User < ApplicationRecord
    def say(text)
      puts text
    end

    def sum(x, y)
      x + y
    end
  end
RUBY

enum = Enumerator.new { |y|
  nodes = ast.children
  until nodes.empty?
    nodes.select! { _1.is_a?(RubyVM::AbstractSyntaxTree::Node) }
    nodes.each { y << _1 }
    nodes = nodes.flat_map(&:children)
  end
}

method_names = enum.select { _1.type == :DEFN }.map { _1.children[0] }
pp method_names
#=> [:say, :sum]