Ruby のメソッドのブロックのソースを取り出す

概要

Ruby のメソッドのブロックのソースを取り出すコードを書いてみました。

block_source.rb

def inspector(ast:, &block)
  ast.children.each do |child|
    next unless child.instance_of?(RubyVM::AbstractSyntaxTree::Node)
    yield child
    inspector(ast: child, &block)
  end
end

def find_node(ast:, type:, lineno:)
  inspector(ast: ast) do |node|
    return node if node.type == type && node.first_lineno == lineno
  end

  nil
end

def extract_source(node:, source:)
  first_lineno = node.first_lineno - 1
  first_column = node.first_column
  last_lineno = node.last_lineno - 1
  last_column = node.last_column - 1

  if first_lineno == last_lineno
    source[first_lineno][first_column..last_column]
  else
    src = ' ' * first_column + source[first_lineno][first_column..]
    ((first_lineno + 1)...last_lineno).each do |lineno|
      src << source[lineno]
    end
    src << source[last_lineno][0..last_column]
  end
end

module Kernel
  # RUBY_VERSION >= '2.6.0'
  def block_source
    @file_specs ||= {}
    bl = caller_locations.last
    source = @file_specs.dig(bl.path, :source) || File.readlines(bl.path)
    @file_specs[bl.path] ||= { source: source, ast: RubyVM::AbstractSyntaxTree.parse(source.join) }
    node = find_node(
      ast: @file_specs.dig(bl.path, :ast),
      type: :ITER,
      lineno: bl.lineno
    )
    extract_source(node: node.children[1], source: source) if node
  end
end

使用法

foo.rb

require_relative 'block_source'

def foo
  pp block_source
end

foo { 'hello' }   #=> " { 'hello' }"
foo { |i| i * 3 } #=> " { |i| i * 3 }"
foo               #=> nil

コード解説

まず bl = caller_locations.last で block_source の呼び出し元の情報を取得。

bl = caller_locations.last

次に find_node ブロック付きのメソッド呼び出しである ITER ノードを探します。

    node = find_node(
      ast: @file_specs.dig(bl.path, :ast),
      type: :ITER,
      lineno: bl.lineno
    )

find_node の実装自体はかなりシンプルで inspector ですべてのノードの中から指定した type と lineno に一致したものを返すようにしています。

def find_node(ast:, type:, lineno:)
  inspector(ast: ast) do |node|
    return node if node.type == type && node.first_lineno == lineno
  end

  nil
end

最後に見つけたノードからブロック部分のソースコードを抜き出します。 ITER ノードは children の 2番目の要素に block の情報が入っているのでそれを渡します。

extract_source(node: node.children[1], source: source) if node

TODO

個人的にあまりそういうコードを書かないので無視しましたが、同じ行でブロック付きメソッド呼び出しを複数回行うと、全部最初のブロックのソースしか取ってこないので、ここはもうちょっと頑張る必要がありそうです。（呼び出し元の column 情報が取れると良さそう

foo { 'hello' }; bar { 'world' }

" { 'hello' }"
" { 'hello' }"