Ruby で簡単に既存メソッドの前後に処理を挟み込む方法

自分用の備忘録。このコード、ブログに書いておかないと忘れそうなので…

やりたかったこと

RubyKaigi の型の話を聞いて、「メソッドの実行を上書きして、引数の型を読み取る」ってことができないかなと思いついて実装したコード。

コード

既存の出力を { ... } と丸括弧で囲むようにするコード。

class User
  def self.a
    print 'singleton method'
  end

  def a
    print 'a'
  end

  def b(a1)
    print a1
  end

  def c(a1, a2)
    print [a1, a2]
  end

  def d(a, b, c = {})
    print [a, b, c]
  end

  def e(&_)
    yield
  end

  def f
    yield
  end
end

module M
  class << self
    def call_original(obj, name, *args, &block)
      m = obj.method(name)
      params = m.parameters
      has_yield = block && params.all? { |type, _name| type != :block }

      if params.empty?
        if has_yield
          m.call { block.call }
        else
          m.call
        end
      elsif has_yield
        m.call(*args) { block.call }
      else
        m.call(*args, &block)
      end
    end

    def method_names_on(obj)
      obj.public_instance_methods - obj.superclass.public_instance_methods
    end

    def wrap_method(obj, name)
      obj.send(:define_method, name) do |*args, &block|
        print '{ '
        M.call_original(self, name, *args, &block)
        puts ' }'
      end
    end

    def spy(klass)
      Module.new do
        refine klass do
          M.method_names_on(klass).each do |name|
            M.wrap_method(self, name)
          end
        end

        refine klass.singleton_class do
          M.method_names_on(klass.singleton_class).each do |name|
            M.wrap_method(self, name)
          end
        end
      end
    end
  end
end

using M.spy(User)
User.a
#=> { singleton method }
user = User.new
user.a
#=> { a }
user.b(1)
#=> { 1 }
user.c(1, 2)
#=> { [1, 2] }
user.d(1, 2, a: 3)
#=> { [1, 2, {:a=>3}] }
user.e { print 'e' }
#=> { e }
user.f { print 'f' }
#=> { f }