Lua で with-output-to-string

「Gauche クックブック」というサイトの記事 [1] のコード例で初めて知って面白いと思ったのが、Gauche には with-output-to-string というフォームがあって（Scheme 全般にあるのかとか Common Lisp にもあるのかは調べていない）、これでコードを囲むだけでその中身のコードで標準出力に出力しようとした文字列を全部バッファに溜めて文字列として返してくれるらしい。
要は一種のストリングバッファなのだけど、こういう書き方ができるのは自然ですばらしいと思った。これを Lua でやりたい。

Lua の print や io.stdout:write はそのままでは単に標準出力に文字を出すしかできないので、これらを一時的に取り替えて代わりにバッファに溜めるようにする。高階関数としても実装できるのかもしれないけど Metalua のマクロを使ってみた。

require "std"

-{block:

  function with_output_to_string_builder(x)
    local block = unpack(x)
    return +{
      function ()
        local buffer = {} -- string buffer

        -- back up original print and io
        local _print = print
        local _io = io

        -- our new print
        print = function(...)
          buffer[#buffer+1] = table.concat(table.imap(|x| tostring(x), {...}), "\t")
          buffer[#buffer+1] = "\n"
        end

        -- our new io.stdout and io.write
        io = {
          stdout = {
            write = function(self, ...)
              for i,v in ipairs{...} do
                buffer[#buffer+1] = tostring(v)
              end
            end
          },
          write = function(...)
              io.output:write(...)
          end,
        }
        io.output = io.stdout
        setmetatable(io.stdout, {__index = _io.stdout})
        setmetatable(io, {__index = _io})

        -- do the work
        -{block}

        -- restore original print and io
        print = _print
        io = _io

        return table.concat(buffer)
      end ()
    }
  end

  mlp.lexer:add{"with_output_to_string"}
  mlp.expr:add{
    "with_output_to_string", mlp.block, "end",
    builder = with_output_to_string_builder
  }
}

これを使うと以下のような書き方で、

local s = with_output_to_string
  print("hello", "world")

  io.write("howdy, ", "world\n")

  io.output = io.open("/dev/null")
  io.write("this message will be discarded")

  io.output = io.stdout
  io.write("howdy, ", "again\n")

  io.stdout:write("good", "bye\n")
end

変数 s には "hello\tworld\nhowdy, world\nhowdy, again\ngoodbye\n" が入る。

文字列を繰り返す関数は（Lua には string.rep があるが）以下のように書ける。

function myrep(str, n)
  return with_output_to_string
    for i=1,n do
      io.write(str)
    end
  end
end

グローバルな print 関数などを直接交換するので既存の関数の中で print などを使っていてもバッファに入る。

function f()
  print("foo")
end

print(string.upper(with_output_to_string f() end)) --> FOO

ただし C による拡張の中で標準出力に出しているようなものは当然入らないだろう。
あとコルーチンの中で with_output_to_string を使ってその中で yield したような場合不都合が起きそう。

[1] http://d.hatena.ne.jp/rui314/20070416/p1