ruby unit tests refactored

[biopieces.git] / code_ruby / Maasha / lib / biopieces.rb

require 'optparse'
require 'open3'
require 'pp'

# Error class for all exceptions to do with option casts.
class CastError < StandardError
end

# Biopieces are command line scripts and uses OptionParser to parse command line
# options according to a list of casts. Each cast prescribes the long and short
# name of the option, the type, if it is mandatory, the default value, and allowed
# and disallowed values. An optional list of extra casts can be supplied, and the
# integrity of the casts are checked. Following the command line parsing, the
# options are checked according to the casts. Methods are also included for handling
# the parsing and emitting of Biopiece records, which are ASCII text records consisting
# of lines with a key/value pair seperated by a colon and a white space ': '.
# Each record is separated by a line with three dashes '---'.
class Biopieces
  TYPES        = %w[flag string list int uint float file file! files files! dir dir! genome]
  MANDATORY    = %w[long short type mandatory default allowed disallowed]
  REGEX_LIST   = /^(list|files|files!)$/
  REGEX_INT    = /^(int|uint)$/
  REGEX_STRING = /^(string|file|file!|dir|dir!|genome)$/

  # Initialize a Biopiece and write the status to file.
  def initialize(no_status=nil)
    status_set unless no_status
  end

  # Check the integrity of a list of casts, followed by parsion options from argv
  # and finally checking the options according to the casts. Returns nil if
  # argv is empty, otherwise an options hash.
  def parse(argv,casts=[],script_path=$0)
    @casts       = casts
    @script_path = script_path

    cast_ubiquitous
    cast_check

    @options = {}

    options_template = OptionParser.new do |option|
      @casts.each do |cast|
        if cast[:type] == 'flag'
          option.on("-#{cast[:short]}", "--#{cast[:long]}") do |o|
            @options[cast[:long]] = o
          end
        elsif cast[:type] =~ REGEX_LIST
          option.on( "-#{cast[:short]}", "--#{cast[:long]} A", Array) do |a|
            @options[cast[:long]] = a
          end
        elsif cast[:type] =~ REGEX_INT
          option.on("-#{cast[:short]}", "--#{cast[:long]} I", Integer) do |i|
            @options[cast[:long]] = i
          end
        elsif cast[:type] =~ REGEX_STRING
          option.on("-#{cast[:short]}", "--#{cast[:long]} S", String) do |s|
            @options[cast[:long]] = s
          end
        elsif cast[:type] == 'float'
          option.on("-#{cast[:short]}", "--#{cast[:long]} F", Float) do |f|
            @options[cast[:long]] = f
          end
        else
          raise ArgumentError, "Unknown option type: '#{cast[:type]}'"
        end
      end
    end

    options_template.parse!(argv)

    if print_usage_full?
      print_usage_and_exit(true)
      return # break for unit testing.
    elsif print_usage_short?
      print_usage_and_exit
      return # break for unit testing.
    end

    options_default
    options_glob
    options_check

    @options
  end

  # Open Biopiece input stream if not open and iterate over all Biopiece
  # records in the stream.
  def each_record
    @in = stream_in_open unless @in.is_a? IO

    record = {}

    @in.each_line do |line|
      case line
      when /^([^:]+): (.*)$/
        record[$1] = $2
      when /^---$/
        yield record unless record.empty?
        record = {}
      else
        raise "Bad record format: #{line}"
      end
    end

    yield record unless record.empty?

    self # conventionally
  end

  alias :each :each_record

  # Open Biopiece output stream if not open and puts record to the stream.
  def puts(record)
    @out = stream_out_open unless @out.is_a? IO

    record.each do |key,value|
      @out.print "#{key}: #{value}\n"
    end

    @out.print "---\n"
  end

  # Close Biopiece streams, remove tmp_dir, end log status.
  def clean
    @in.close  if @in.respond_to?  :close
    @out.close if @out.respond_to? :close
    # remove tmpdir if found
    status_log
    # remove status file
  end

  private

  # Given the script name determine the path of the wiki file with the usage info.
  def wiki_path
    path = ENV["BP_DIR"] + "/bp_usage/" + File.basename(@script_path, ".rb") + ".wiki"
    raise "No such wiki file: #{path}" unless File.file? path

    path
  end

  # Add ubiquitous options casts.
  def cast_ubiquitous
    @casts << {:long => 'help',       :short => '?', :type => 'flag',   :mandatory => false, :default => nil, :allowed => nil, :disallowed => nil}
    @casts << {:long => 'stream_in',  :short => 'I', :type => 'files!', :mandatory => false, :default => nil, :allowed => nil, :disallowed => nil}
    @casts << {:long => 'stream_out', :short => 'O', :type => 'file',   :mandatory => false, :default => nil, :allowed => nil, :disallowed => nil}
    @casts << {:long => 'verbose',    :short => 'v', :type => 'flag',   :mandatory => false, :default => nil, :allowed => nil, :disallowed => nil}
  end

  # Check integrity of the casts.
  def cast_check
    cast_check_keys
    cast_check_values
    cast_check_duplicates
  end
  
  # Check if all mandatory keys are present in casts and raise if not.
  def cast_check_keys
    @casts.each do |cast|
      MANDATORY.each do |mandatory|
        raise CastError, "Missing symbol in cast: '#{mandatory.to_sym}'" unless cast.has_key? mandatory.to_sym
      end
    end
  end

  # Check if all values in casts are valid.
  def cast_check_values
    @casts.each do |cast|
      cast_check_val_long(cast)
      cast_check_val_short(cast)
      cast_check_val_type(cast)
      cast_check_val_mandatory(cast)
      cast_check_val_default(cast)
      cast_check_val_allowed(cast)
      cast_check_val_disallowed(cast)
    end
  end

  # Check if the values to long are legal and raise if not.
  def cast_check_val_long(cast)
    unless cast[:long].is_a? String and cast[:long].length > 1
      raise CastError, "Illegal cast of long: '#{cast[:long]}'"
    end
  end
  
  # Check if the values to short are legal and raise if not.
  def cast_check_val_short(cast)
    unless cast[:short].is_a? String and cast[:short].length == 1
      raise CastError, "Illegal cast of short: '#{cast[:short]}'"
    end
  end

  # Check if values to type are legal and raise if not.
  def cast_check_val_type(cast)
    type_hash = {}
    TYPES.each do |type|
      type_hash[type] = true
    end

    unless type_hash.has_key? cast[:type]
      raise CastError, "Illegal cast of type: '#{cast[:type]}'"
    end
  end

  # Check if values to mandatory are legal and raise if not.
  def cast_check_val_mandatory(cast)
    unless cast[:mandatory] == true or cast[:mandatory] == false
      raise CastError, "Illegal cast of mandatory: '#{cast[:mandatory]}'"
    end
  end

  # Check if values to default are legal and raise if not.
  def cast_check_val_default(cast)
    unless cast[:default].nil? or
           cast[:default].is_a? String or
           cast[:default].is_a? Integer or
           cast[:default].is_a? Float
      raise CastError, "Illegal cast of default: '#{cast[:default]}'"
    end
  end

  # Check if values to allowed are legal and raise if not.
  def cast_check_val_allowed(cast)
    unless cast[:allowed].is_a? String or cast[:allowed].nil?
      raise CastError, "Illegal cast of allowed: '#{cast[:allowed]}'"
    end
  end

  # Check if values to disallowed are legal and raise if not.
  def cast_check_val_disallowed(cast)
    unless cast[:disallowed].is_a? String or cast[:disallowed].nil?
      raise CastError, "Illegal cast of disallowed: '#{cast[:disallowed]}'"
    end
  end

  # Check cast for duplicate long or short options names.
  def cast_check_duplicates
    check_hash = {}
    @casts.each do |cast|
      raise CastError, "Duplicate argument: '--#{cast[:long]}'" if check_hash.has_key? cast[:long]
      raise CastError, "Duplicate argument: '-#{cast[:short]}'" if check_hash.has_key? cast[:short]
      check_hash[cast[:long]]  = true
      check_hash[cast[:short]] = true
    end
  end

  # Check if full "usage info" should be printed.
  def print_usage_full?
    @options["help"]
  end

  # Check if short "usage info" should be printed.
  def print_usage_short?
    if not $stdin.tty?
      return false
    elsif @options["stream_in"]
      return false
    elsif @options["data_in"]
      return false
    elsif wiki_path =~ /^(list_biopieces|list_genomes|list_mysql_databases|biostat)$/  # TODO get rid of this!
      return false
    else
      return true
    end
  end

  # Print usage info by Calling an external script 'print_wiki'
  # using a system() call and exit. An optional 'full' flag
  # outputs the full usage info.
  def print_usage_and_exit(full=nil)
    if full
      system("print_wiki --data_in #{wiki_path} --help")
    else
      system("print_wiki --data_in #{wiki_path}")
    end

    raise "Failed printing wiki: #{wiki_path}" unless $?.success?

    exit
  end

  # Set default options value from cast unless a value is set.
  def options_default
    @casts.each do |cast|
      if cast[:default]
        @options[cast[:long]] = cast[:default] unless @options.has_key? cast[:long]
      end
    end
  end

  # Expands glob expressions to a full list of paths.
  # Examples: "*.fna" or "foo.fna,*.fna" or "foo.fna,/bar/*.fna"
  def options_glob
    @casts.each do |cast|
      if cast[:type] == 'files' or cast[:type] == 'files!'
        if @options.has_key? cast[:long]
          files = []
        
          @options[cast[:long]].each do |path|
            if path.include? "*"
              Dir.glob(path).each do |file|
                files << file if File.file? file
              end
            else
              files << path
            end
          end

          @options[cast[:long]] = files
        end
      end
    end
  end

  # Check all options according to casts.
  def options_check
    @casts.each do |cast|
      options_check_mandatory(cast)
      options_check_int(cast)
      options_check_uint(cast)
      options_check_file(cast)
      options_check_files(cast)
      options_check_dir(cast)
      options_check_allowed(cast)
      options_check_disallowed(cast)
    end
  end
  
  # Check if a mandatory option is set and raise if it isn't.
  def options_check_mandatory(cast)
    if cast[:mandatory]
      raise ArgumentError, "Mandatory argument: --#{cast[:long]}" unless @options.has_key? cast[:long]
    end
  end

  # Check int type option and raise if not an integer.
  def options_check_int(cast)
    if cast[:type] == 'int' and @options.has_key? cast[:long]
      unless @options[cast[:long]].is_a? Integer
        raise ArgumentError, "Argument to --#{cast[:long]} must be an integer, not '#{@options[cast[:long]]}'"
      end
    end
  end
  
  # Check uint type option and raise if not an unsinged integer.
  def options_check_uint(cast)
    if cast[:type] == 'uint' and @options.has_key? cast[:long]
      unless @options[cast[:long]].is_a? Integer and @options[cast[:long]] >= 0
        raise ArgumentError, "Argument to --#{cast[:long]} must be an unsigned integer, not '#{@options[cast[:long]]}'"
      end
    end
  end

  # Check file! type argument and raise if file don't exists.
  def options_check_file(cast)
    if cast[:type] == 'file!' and @options.has_key? cast[:long]
      raise ArgumentError, "No such file: '#{@options[cast[:long]]}'" unless File.file? @options[cast[:long]]
    end
  end

  # Check files! type argument and raise if files don't exists.
  def options_check_files(cast)
    if cast[:type] == 'files!' and @options.has_key? cast[:long]
      @options[cast[:long]].each do |path|
        raise ArgumentError, "No such file: '#{path}'" unless File.file? path
      end
    end
  end
  
  # Check dir! type argument and raise if directory don't exist.
  def options_check_dir(cast)
    if cast[:type] == 'dir!' and @options.has_key? cast[:long]
      raise ArgumentError, "No such directory: '#{@options[cast[:long]]}'" unless File.directory? @options[cast[:long]]
    end
  end
  
  # Check options and raise unless allowed.
  def options_check_allowed(cast)
    if cast[:allowed] and @options.has_key? cast[:long]
      allowed_hash = {}
      cast[:allowed].split(',').each { |a| allowed_hash[a] = 1 }
  
      raise ArgumentError, "Argument '#{@options[cast[:long]]}' to --#{cast[:long]} not allowed" unless allowed_hash.has_key? @options[cast[:long]]
    end
  end
  
  # Check disallowed argument values and raise if disallowed.
  def options_check_disallowed(cast)
    if cast[:disallowed] and @options.has_key? cast[:long]
      cast[:disallowed].split(',').each do |val|
        raise ArgumentError, "Argument '#{@options[cast[:long]]}' to --#{cast[:long]} is disallowed" if val == @options[cast[:long]]
      end
    end
  end

  # Open Biopieces input data stream for reading from either
  # stdin or from a list of files specified in options["stream_in"].
  def stream_in_open
    if not $stdin.tty?
      stream = $stdin
    else
      stream = read(@options["stream_in"])
    end

    stream
  end

  # Open Biopieces output data stream for writing to stdout
  # or a file specified in options["stream_out"].
  def stream_out_open
    if @options["stream_out"]
      stream = write(@options["stream_out"], @options["compress"])
    else
      stream = $stdout
    end

    stream
  end

  # Opens a reads stream to a list of files.
  def read(files)
    if zipped?(files)
      stream = zread(files)
    else
      stream = nread(files)
    end

    stream
  end

  # Opens a write stream to a file and returns a _io_ object.
  def write(file, zip=nil)
    zip ? zwrite(file) : nwrite(file)
  end

  # Test if a list of files are gzipped or not.
  # Raises if files are mixed zipped and unzipped.
  def zipped?(files)
    type_hash = {}

    files.each do |file|
      type = `file #{file}`

      if type =~ /gzip compressed/
        type_hash[:gzip] = true
      else
        type_hash[:ascii] = true
      end
    end

    raise "Mixture of zipped and unzipped files" if type_hash.size == 2

    type_hash[:gzip]
  end

  # Opens a list of gzipped files for reading and return an _io_ object.
  def zread(files)
    stdin, stdout, stderr = Open3.popen3("zcat " + files.join(' '));
    stdin.close
    stderr.close
    stdout
  end

  # Opens a file for gzipped writing and return an _io_ object.
  def zwrite(file)
    stdin, stdout, stderr = Open3.popen3("gzip -f > #{file}")
    stderr.close
    stdout.close
    stdin
  end

  # Opens a list of files for reading and return an _io_ object.
  def nread(files)
    stdin, stdout, stderr = Open3.popen3("cat " + files.join(' '));
    stdin.close
    stderr.close
    stdout
  end

  # Opens a file for writing and return an _io_ object.
  def nwrite(file)
    File.open(file, mode="w")
  end

  # Write the status to a status file.
  def status_set
    now    = Time.new
    time   = now.strftime("%Y-%m-%d %X")
    user   = ENV["USER"]
    script = File.basename($0, ".rb")
    pid    = $$
    path   = ENV["BP_TMP"] + "/" + [user, script, pid, "status"].join(".")

    File.open(path, mode="w") { |file| file.puts [time, ARGV.join(" ")].join(";") }
  end

  # Write the Biopiece status to the log file.
  def status_log(status="OK")
    now     = Time.new
    time1   = now.strftime("%Y-%m-%d %X")
    user    = ENV["USER"]
    script  = File.basename($0, ".rb")
    pid     = $$
    path    = ENV["BP_TMP"] + "/" + [user, script, pid, "status"].join(".")

    stream = File.open(path)
    time0, args, tmp_dir = stream.first.split(";")

    # Dir.rmdir(tmp_dir) unless tmp_dir.nil? and File.directory? tmp_dir   #TODO fix this!

    elap     = time_diff(time0, time1)
    command  = [script, args].join(" ") 
    log_file = ENV["BP_LOG"] + "/biopieces.log"

    File.open(log_file, mode="a") { |file| file.puts [time0, time1, elap, user, status, command].join("\t") }
  end

  # Get the elapsed time from the difference between two time stamps.
  def time_diff(t0, t1)
    t0 =~ /(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)/
    year0 = $1.to_i
    mon0  = $2.to_i
    day0  = $3.to_i
    hour0 = $4.to_i
    min0  = $5.to_i
    sec0  = $6.to_i

    sec0 += day0  * 24 * 60 * 60
    sec0 += hour0 * 60 * 60
    sec0 += min0  * 60

    t1 =~ /(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)/
    year1 = $1.to_i
    mon1  = $2.to_i
    day1  = $3.to_i
    hour1 = $4.to_i
    min1  = $5.to_i
    sec1  = $6.to_i

    sec1 += day1  * 24 * 60 * 60
    sec1 += hour1 * 60 * 60
    sec1 += min1  * 60

    year = year1 - year0
    mon  = mon1  - mon0
    day  = day1  - day0 

    sec  = sec1 - sec0

    hour = ( sec / ( 60 * 60 ) ).to_i
    sec -= hour * 60 * 60

    min  = ( sec / 60 ).to_i
    sec -= min * 60

    [sprintf("%02d", hour), sprintf("%02d", min), sprintf("%02d", sec)].join(":")
  end
end

__END__