Test script for learning about the PTY (pseudoterminal) lib

pty-test.rb

#!/usr/bin/env ruby
# -*- coding: utf-8; -*-
#
# PTY test program
# Séan O'Halpin, 2012-06-29
#
# examples:
#
#   ./pty-test.rb irb --errpipe --childraw
#
#   ./pty-test.rb vi --raw
#
# note that this leaves the host terminal at 80x25
#
require 'pty'
require 'time'
require 'io/console' # for IO#raw!

def dbg(*a)
  if $OPT_DEBUG
    File.open("dbg.log", "a:utf-8") do |file|
      file.puts Time.now.xmlschema + ": " + a.inspect
    end
  end
end

def wrap_text(text, wrap_txt = "\\1")
  text.gsub(/^(.*)$/m, wrap_txt)
end

# quick and dirty option handling
$OPT_CHILDRAW = ARGV.delete("--childraw") # set slave to raw
$OPT_DEBUG    = ARGV.delete("--debug")    # output debug info into dbg.log
$OPT_ERRPIPE  = ARGV.delete("--errpipe")  # use pipe for child's STDERR
$OPT_NOECHO   = ARGV.delete("--noecho")   # set slave to noecho
$OPT_RAW      = ARGV.delete("--raw")      # set parent terminal to raw
$OPT_UPCASE   = ARGV.delete("--upcase")   # for example - will interfere with raw output

# number of chars to read at a time
NCHARS = 4096

if ARGV.size > 0
  cmd = ARGV.join(" ")
else
  # example command
  cmd = "irb --simple-prompt"
end

begin
  if $OPT_ERRPIPE
    # create pipe to handle STDERR from child separately from STDOUT
    errpipe_read, errpipe_write = IO.pipe
  else
    errpipe_read = errpipe_write = :INVALID_FD # won't match in case below
  end

  PTY.open do |master, slave|
    if pid = fork
      # parent

      # close slave
      slave.close
      # close the write end of the error pipe
      errpipe_write.close if $OPT_ERRPIPE

      # optionally set our own STDOUT to raw for child programs which
      # themselves put the terminal in raw mode, e.g. vi
      if $OPT_RAW
        STDOUT.raw!
      end

      # handle i/o using a select loop
      res = catch :eof do
        loop do
          if $OPT_ERRPIPE
            #              [read, write, exception]
            rv = IO.select [errpipe_read, master, STDIN], [], [errpipe_read, master, STDIN]
          else
            rv = IO.select [master, STDIN], [], [master, STDIN]
          end

          if rv
            ra, wa, ea = *rv

            if ra.size > 0
              ra.each do |rfd|
                case rfd

                when STDIN
                  begin
                    # write to STDIN of child process
                    text = rfd.readpartial(NCHARS)
                    dbg :STDIN, text
                    # convert F5 into save and quit in vi
                    text = text.gsub(/\e\[15~/, "\e:wq\n")
                    master.print text
                  rescue EOFError => e
                    throw :eof, "EOF on STDIN"
                  end

                when errpipe_read
                  begin
                    text = rfd.readpartial(NCHARS)
                    dbg :ERROR, text
                    # output text in red
                    print wrap_text(text, "\e[0;31m\\1\e[0m")
                  rescue EOFError
                    throw :eof, "EOF on ERROR"
                  rescue Errno::EIO
                    throw :eof, "EIO on ERROR"
                  end

                when master
                  # read child's STDOUT
                  begin
                    text = rfd.readpartial(NCHARS)
                    dbg :CHILD_STDOUT, text
                    if $OPT_UPCASE
                      text = text.upcase
                    end
                    print text
                  rescue EOFError
                    throw :eof, "EOF on master"
                  rescue Errno::EIO
                    throw :eof, "EIO on master"
                  end

                else
                  throw :eof, "Unknown fd"
                end

              end
            end

            # haven't seen this - not sure under what circumstances
            # this will arise - need to RTFM
            if ea.size > 0
              throw :eof, "select returned exception - #{ea}"
            end
          else
            throw :eof, "select returned nil"
          end
        end
      end

      # reset output to normal
      STDOUT.cooked!
      puts "\n#{res}"
      # close master to generate an EOF on child's STDIN (otherwise
      # child won't exit and parent will hang in waitpid)
      master.close
      errpipe_read.close if $OPT_ERRPIPE
      Process.waitpid(pid)
    else
      # child

      # close master (if we don't then hangs on EOF in parent)
      master.close
      # redirect STDERR to pipe
      if $OPT_ERRPIPE
        STDERR.reopen(errpipe_write)
      else
        STDERR.reopen(slave)
      end

      # optionally disable echo and newline conversion
      if $OPT_CHILDRAW
        slave.raw!
      end
      if $OPT_NOECHO
        slave.echo = false
      end

      # redirect STDIN and STDOUT to slave
      STDIN.reopen(slave)
      STDOUT.reopen(slave)
      # don't need slave handle any more
      slave.close

      # replace child process with executed cmd
      begin
        exec cmd
      rescue => e
        STDERR.puts "\r#{e}"
      end
    end
  end
ensure
  # make sure we restore the terminal to its (presumed) original state
  STDOUT.cooked!
end