RubyでYARD定義を使って実行時にメソッド引数と戻り値の型チェックを試みる

Ruby のカンファレンスの度に、毎度話題になる?) Ruby に型が欲しい件ですが、個人的な見解を書いておこうと思います。ちなみに私は RubyKaigi 2018 に参加しておりません。Twitterのタイムラインでの賑わいを見ていただけです。

最近、Swift を触っててそのコンパイルの重さにうんざりしているので、 型推論 は現代のマシンスペックでは基本的に辛いと思っています(カフェでノマドコーディングしたいので)。またメタプログラミングし放題の Ruby に導入するのは困難という認識です。

では、どういうときに 型定義 が欲しくなるのか考えてみます。

  1. コードを書いて実行して確認というトライアンドエラーを減らしたい、実行前にエラーをなるべく洗い出したい。
  2. 型定義起因による補完を効かせながらコーディングしたい。これはエディタやIDEのサポートも必要です。

1 について、Rubyの場合テストでカバレッジを稼いで、なるべくエラーの芽を潰しておこうとします。しかしユニットテストにおいて end-to-end テストが存在しない場合、単純なユーティリティ関数でなければ、モックやスタブに頼るようになりチェックが甘くなります。そのためテストは通過するものの、実際通しで動かしたとき想定とは異なる型のデータ受渡しが発生してしまうことがあります。

2 について、メソッド名や引数名から型をコードの読解で推測することは可能ですが、それなりの規模のアプリケーションやライブラリではコードコメントでドキュメント定義していないと(昔の自分や)他人の書いたコードを扱うのが困難になることが多いと思います。 そして例えば YARD で引数や戻り値の型をコメントに定義して、ドキュメントを生成したりIDEでコード補完に用いることが多いでしょう。

さてここから本題ですが、1 で型チェックしないで検証から漏れてしまう問題は、テストでメソッド呼び出しの引数と戻り値をよりチェックするアサーションを明示的に書いていく必要があります。一方 2 でYARDによるドキュメントの型定義は必ずしも実際のソースコードで走る処理と一致してるか保証されない問題もあります。

そこで『YARD定義によるメソッド引数と戻り値の型チェック』を実行時に行ってみたらどうかと考えてみました。

YARDはテンプレートで書きだす直前の解析データを、YARD::Registry から参照する機能を提供しています。またあらゆるメソッド呼び出し(:call)と戻り(:return)は、TracePoint 機構を使ってフックすることができます。これを組み合わせれば、割と簡単に、実行時に定義どおりに適切な型の引数が指定されてメソッド呼び出され、適切な戻り値が返っているかをチェックすることできると思います。

ということで、試してみましょう。

lib/dog.rb

これが YARD 定義を書いた検証対象のクラスになります。

module Animal
  #
  # This class is Dog
  #
  class Dog
    # @param name [String] a name
    # @param weight [Numeric] weight
    def initialize(name, weight)
      @name = name
      @weight = weight
      @children = []
    end

    # Add a child dog
    # 
    # @param dog [Animal::Dog] a child dog
    def add_child(dog)
      @children.push(dog)
    end

    # Run.
    # 
    # @param distance [Integer]
    # @return [String] message
    def run(distance)
      "#{@name} runs #{distance}."
    end

    # dummy method returns wrong type value
    # 
    # @return [Integer]
    def dummy
      "a string"
    end
  end
end

definition.rb

YARD::Registry から ClassObject と属する MethodObject を解析してチェックしやすい定義クラス MethodDefinition の集合 DefinitionStore に変換します。

require 'rubygems'
require 'yard'

class MethodDefinition
  def initialize(method_obj)
    @name = method_obj.name(true)
    @args = []
    @ret = nil

    load_docstr(method_obj.docstring)
  end

  def validate_arguments(args)
    errors = []

    args.each_with_index do |arg, i|
      arg_def = @args[i]

      ts = arg_def.types
      result = ts.find { |t|
          klass = Object.const_get(t)
          arg.is_a?(klass)
        }
      
      if result.nil?
        if ts.count > 1
          errors.push "#{arg_def.name}: #{arg.inspect} isn't any of " + ts.join(',')
        else
          errors.push "#{arg_def.name}: #{arg.inspect} isn't #{ts[0]}"
        end
      end
    end

    errors.empty? ? nil : "#{@name}(" + errors.join(', ') + ")"
  end

  def validate_return(ret_val)
    return nil if @ret.nil?

    ts = @ret.types
    result = ts.find { |t|
        klass = Object.const_get(t)
        ret_val.is_a?(klass)
      }

    if result.nil?
      if ts.count > 1
        "#{@name} returned #{ret_val.inspect} isn't any of " + ts.join(',') + ")"
      else
        "#{@name} returned #{ret_val.inspect} isn't #{ts[0]})"
      end
    else
      nil
    end
  end

  def load_docstr(docstr)
    docstr.tags.each do |tag|
      tag_name = tag.tag_name

      if tag_name == 'param'
        @args.push OpenStruct.new({
            name: tag.name,
            types: tag.types            
          })
      elsif tag_name == 'return'
        @ret = OpenStruct.new({ types: tag.types })
      end
    end
  end
end

class DefinitionStore
  def initialize
    @store = {}

    registry = YARD::Registry.load!
    registry.all(:class).each do |class_obj|
      add(class_obj)
    end
  end

  def add(class_obj)
    method_map = @store[class_obj.path] = {}
    class_obj.meths.each do |mt|
      method_map[mt.name] = MethodDefinition.new(mt)
    end
  end

  def get(klass_name, method_name)
    klass_def = @store[klass_name]
    klass_def[method_name]
  end
end

checker.rb

TracePoint でメソッド呼び出しと戻りをトレースして、引数または戻り値がYARD定義と適合してるかチェックします。

require 'tracer'
require_relative 'definition'

definition_store = DefinitionStore.new

TracePoint.trace(:call, :return) do |tp|
  klass_name = tp.defined_class.name
  method_def = definition_store.get(klass_name, tp.method_id)
  
  if tp.event == :call
    args = tp.binding.local_variables.map do |name|
             tp.binding.local_variable_get(name)
           end

    if err = method_def.validate_arguments(args)
      puts "Invalid call #{klass_name}#{err} on #{tp.path}:#{tp.lineno}"
    end
  elsif tp.event == :return && tp.method_id != :initialize
    if err = method_def.validate_return(tp.return_value)
      puts "Invalid return #{klass_name}#{err} on #{tp.path}:#{tp.lineno}"
    end
  end
end

# 試行
require_relative 'lib/dog'

dog1 = Animal::Dog.new(nil, "4.5")
p dog1.run(20)
dog1.dummy

dog2 = Animal::Dog.new("Taro", 6.5)
dog2.add_child(dog1)
dog2.add_child(nil)
p dog2.run(nil)

実行例

yard で解析データを生成してから、下記のようにチェッカーを実行してみます。

$ ruby cheker.rb
Invalid call Animal::Dog#initialize(name: nil isn't String, weight: "4.5" isn't Numeric) on ./lib/dog.rb:8
" runs 20."
Invalid return Animal::Dog#dummy returned "a string" isn't Integer) on ./lib/dog.rb:34
Invalid call Animal::Dog#add_child(dog: nil isn't Animal::Dog) on ./lib/dog.rb:17
Invalid call Animal::Dog#run(distance: nil isn't Integer) on ./lib/dog.rb:25
"Taro runs ."

とりあえずこの簡単な例で試すことは成功しました。もちろん引数が特定のモジュールを mix-in してるかといったチェックなど全然足りてはいません。ただこのような機構を、テストや開発向けデプロイ環境で動かすことで型違いを起因とするバグを減らせる手法の1つとして、提示できたんではないかと思います。