Ruby
YAML

RubyでYAMLを拡張する方法と、外部ファイルのincludeに対応させる実例

More than 1 year has passed since last update.

設定ファイルとしてYAMLを扱うことが多いため、yamlを便利に扱うためにyaml_vaultyaml_masterというgemを作っている。

yaml_vaultはYAMLの特定キーを暗号化・複合化するためのgemで、yaml_masterはmasterファイルとして指定されたyamlから一部のキーだけを切り出して別ファイルに書き出したり、まとめて設定ファイルを生成したりするためのgemだ。

そして、最近この二つのgemを改修するためにRubyのYAML生成処理を調べたので、その知見をまとめておく。

yaml_vaultは元々の動作では暗号化したり複合化したりすると、YAMLのエイリアス記法が展開されてしまうという動作上の限界があった。

こういうのが

default: &default
  adapter: mysql
  encoding: utf8mb4
  username root

production
  <<: *default
  password: secretpass

こうなる

default: &default
  adapter: mysql
  encoding: utf8mb4
  username root

production:
  adapter: mysql
  encoding: utf8mb4
  username root
  password: <encrypted string>

また、yaml_masterは外部ファイルをincludeするための仕組みを実装していたのだが、それだとエイリアス記法が共有できないという限界があった。

yaml_master:
  db: db.yml

default: &default
  adapter: mysql
  encoding: utf8mb4
  username root

data
  !include production.yml
# production.yml

production:
  <<: *default
  password: secretpass

こういう書き方をするとアンカーが見つからなくてエラーになる。

この二つの問題を解決するには、YAMLのASTを構築する段階で既に必要な処理が終わってる状態を作る必要がある。

RubyのYAMLライブラリはpsychというバックエンドライブラリが実際の処理をしている。今回、AST構築処理を弄るためにどこに処理を差し込めばいいかを調べた。

psychはPsych::ParserというクラスがYAMLのトークナイズをしてハンドラクラスのイベントハンドラメソッドを呼び出す。例えば処理対象がsequenceだったらstart_sequenceend_sequenceを呼ぶ。
Psych::TreeBuilder < Psych::Handlerというクラスはイベントハンドラメソッドで、Psych::Nodes::<NodeType>クラスのインスタンスを生成し、YAMLのASTを構築する。

Psych::TreeBuilderの実装は以下の様になっている。

# frozen_string_literal: false
require 'psych/handler'

module Psych
  ###
  # This class works in conjunction with Psych::Parser to build an in-memory
  # parse tree that represents a YAML document.
  #
  # == Example
  #
  #   parser = Psych::Parser.new Psych::TreeBuilder.new
  #   parser.parse('--- foo')
  #   tree = parser.handler.root
  #
  # See Psych::Handler for documentation on the event methods used in this
  # class.
  class TreeBuilder < Psych::Handler
    # Returns the root node for the built tree
    attr_reader :root

    # Create a new TreeBuilder instance
    def initialize
      @stack = []
      @last  = nil
      @root  = nil
    end

    %w{
      Sequence
      Mapping
    }.each do |node|
      class_eval %{
        def start_#{node.downcase}(anchor, tag, implicit, style)
          n = Nodes::#{node}.new(anchor, tag, implicit, style)
          @last.children << n
          push n
        end

        def end_#{node.downcase}
          pop
        end
      }
    end

    ###
    # Handles start_document events with +version+, +tag_directives+,
    # and +implicit+ styling.
    #
    # See Psych::Handler#start_document
    def start_document version, tag_directives, implicit
      n = Nodes::Document.new version, tag_directives, implicit
      @last.children << n
      push n
    end

    ###
    # Handles end_document events with +version+, +tag_directives+,
    # and +implicit+ styling.
    #
    # See Psych::Handler#start_document
    def end_document implicit_end = !streaming?
      @last.implicit_end = implicit_end
      pop
    end

    def start_stream encoding
      @root = Nodes::Stream.new(encoding)
      push @root
    end

    def end_stream
      pop
    end

    def scalar value, anchor, tag, plain, quoted, style
      s = Nodes::Scalar.new(value,anchor,tag,plain,quoted,style)
      @last.children << s
      s
    end

    def alias anchor
      @last.children << Nodes::Alias.new(anchor)
    end

    private
    def push value
      @stack.push value
      @last = value
    end

    def pop
      x = @stack.pop
      @last = @stack.last
      x
    end
  end
end

使う時はこんな感じ。

YAML::Parser.new(YAML::TreeBuilder.new).parse(yaml)

処理スタックを積みながら、ノードを生成しているだけなので簡単な作りになっている。
このクラスを継承して独自の処理を挟み込むことで、ASTを構築する段階で任意のフィルター処理を実行できる。

yaml_vaultではこの様な実装をしている。

require 'yaml'

module YamlVault
  class YAMLTreeBuilder < YAML::TreeBuilder
    def initialize(target_paths, cryptor, mode)
      super()

      @path_stack = []
      @target_paths = target_paths
      @cryptor = cryptor
      @mode = mode
    end

    def start_document(*)
      result = super
      @path_stack.push "$"
      result
    end

    def end_document(*)
      @path_stack.pop
      super
    end

    def start_mapping(*)
      if YAML::Nodes::Sequence === @last
        current_path = @last.children.size
        @path_stack << current_path
      end

      super
    end

    def end_mapping(*)
      @path_stack.pop
      super
    end

    def start_sequence(*)
      if YAML::Nodes::Sequence === @last
        current_path = @last.children.size
        @path_stack << current_path
      end

      super
    end

    def end_sequence(*)
      @path_stack.pop
      super
    end

    def scalar(value, anchor, tag, plain, quoted, style)
      result = super

      case @last
      when YAML::Nodes::Sequence
        current_path = @last.children.size - 1
        @path_stack << current_path
      when YAML::Nodes::Mapping
        if @last.children.size.odd?
          @path_stack << value
          return result
        end
      end

      if match_path?
        if @mode == :encrypt
          if tag
            result.value = @cryptor.encrypt("#{tag} #{value}")
            result.tag = nil
            result.plain = true
          else
            result.value = @cryptor.encrypt(value)
          end
        else
          decrypted_value = @cryptor.decrypt(value).to_s
          if decrypted_value =~ /\A(!.*?)\s+(.*)\z/
            result.tag = $1
            result.plain = false
            result.value = $2
          else
            result.value = decrypted_value
          end
        end
      end

      @path_stack.pop

      result
    end

    def alias(anchor)
      @path_stack.pop
      super
    end

    private

    def match_path?
      @target_paths.any? do |target_path|
        target_path.each_with_index.all? do |path, i|
          if path == "*"
            true
          else
            if path.is_a?(Regexp)
              path.match(@path_stack[i])
            else
              path == @path_stack[i]
            end
          end
        end
      end
    end
  end
end

各イベントの処理を上書きして、独自のパススタックを積み、終端であるスカラー値まで到達したら、対象のキーパスを確認して暗号化・複合化の処理を行い、値を差し替える。
この時、valueはまだRubyオブジェクトに変換する前なので文字列のままになっている。

こうやって生成したASTはPsych::Nodes::Stream#to_yamlメソッドを呼ぶことで、ほぼ元のフォーマットを維持したままYAMLとして出力することができる。(流石に空行とかは消えるし、インデントスタイルもちょっと変わる場合がある)

また、YAMLはタグと呼ばれる機能がある。例えば!ruby/range 1..10の様に書くことができる。RubyのYAMLパーサーは、このタグを利用してYAMLが直接サポートしないオブジェクトをYAMLに埋め込むことができる。この例だとパース後にRubyオブジェクトに変換する際に1..10のRangeオブジェクトとして受け取ることができる。
上記実装例のscalarメソッドにはtag="!ruby/range" value="1..10"という値が渡ってくる。

その他、YAMLの書式のバリエーションによって、plain, quoted, styleの値が変わる。どういう書式がどの値になるかは、実際に何かのYAMLをパースさせてppで確認するのが楽だろう。

このタグを利用してyaml_masterの外部ファイルinclude機能も実装できた。

require 'yaml'
require 'pathname'

class YamlMaster::YAMLTreeBuilder < YAML::TreeBuilder
  def initialize(master_path, properties, parser)
    super()
    @master_path = master_path
    @parser = parser
  end

  def scalar(value, anchor, tag, plain, quoted, style)
    case tag
    if "!include"
      ensure_tag_argument(tag, value)

      path = Pathname(value)
      path = path.absolute? ? path : @master_path.dirname.join(path)
      tree = YAML.parse(File.read(path))
      @last.children << tree.children[0]
    else
      super
    end
  end

  private

  def ensure_tag_argument(tag, value)
    if value.empty?
      mark = @parser.mark
      $stderr.puts "tag format error"
      $stderr.puts "#{tag} requires 1 argument at #{mark.line}:#{mark.column}"
      exit 1
    end
  end
end

!includeタグが見つかったら、valueになっている値をファイル名として読み込んだ後YAML.parseをしてインクルード対象のASTを取得する。その後、スカラ値の代わりにそのAST自体を処理中のノードの子として突っ込む。

こうすることで、!include foo.yamlと書いた箇所は、そのままYAMLを記述した場合と完全に同じになる。
結果、ファイルを分割しながら、ベースのファイルにあるエイリアスを共有することができた。
ただし、そうやって書いたファイルは単独では読み込めなくなる。

上記の実装例は見易くするために色々省略をしているが、これを利用すればERBを使わなくても環境変数で値を置き換えたり、外部のテキストファイルから内容を読み込んで差し込む、といった機能が実現できる。

あんまり独自機能をYAMLに足してしまうと、他人に分かりにくくなるので、程々にしておいた方が良いものではあるが、TreeBuilderを弄ることで割と簡単にYAMLのパース結果を弄れるということを覚えておくと、役に立つかもしれない。