RunKit - essentials for scripts and glue

run_kit.rb

# Shared runtime helpers for text/csv/json IO, shell cmds, bannners, etc.

require "csv"
require "digest"
require "json"
require "open3"

module RunKit
  module_function

  #
  # file read/write, including gz
  #

  def file_read(path)
    Pathname(path).then do |path|
      data = path.read
      data = gunzip(data) if path.extname == ".gz"
      data
    end
  end

  def file_write(path, data)
    path = Pathname(path)
    atomic_write(path) do |tmp|
      data = gzip(data) if path.extname == ".gz"
      tmp.write(data)
    end
  end

  #
  # json file read/write, including gz
  #

  def json_read(path, symbolize_names: true)
    JSON.parse(file_read(path), symbolize_names:)
  end

  def json_write(path, json)
    file_write(path, JSON.pretty_generate(json))
  end

  #
  # gzip/gunzip data
  #

  def gzip(str)
    Zlib::GzipWriter.new(StringIO.new).tap do
      _1.write(str)
    end.close.string
  end

  def gunzip(str_gz)
    gz = Zlib::GzipReader.new(StringIO.new(str_gz))
    gz.read
  ensure
    gz&.close
  end

  #
  # CSV read/write
  #

  def csv_read(path, infer: false)
    io = StringIO.new(file_read(path))
    rows = CSV.read(io, encoding: "bom|utf-8")

    headers = rows.shift.map(&:to_sym)
    klass = Struct.new(*headers)

    rows.map do |row|
      row = row.map { _infer_csv(_1) } if infer
      klass.new(*row)
    end
  end

  def csv_write(path, rows, headers: nil)
    atomic_write(path) do |tmp|
      CSV.open(tmp, "wb") { _csv_write0(_1, rows, headers:) }
    end
  end

  def csv_write_stdout(rows, headers: nil)
    CSV($stdout) { _csv_write0(_1, rows, headers:) }
  end

  #
  # shell/shell!
  #

  # Run a command via Open3.capture2e. `cmd` can be passed as:
  #
  # - shell("git status")      # a single string
  # - shell("git", "status")   # varargs strings
  # - shell(["git", "status"]) # an array of strings
  #
  # Prefer arrays so escaping stays explicit and Ruby handles argument
  # boundaries for you. Single strings are convenient but put escaping
  # responsibility on the caller. `vars:` lets you interpolate `{{ hi }}` into
  # the command before it runs. `Pathname` values get shell-escaped, which is
  # real nice here.
  #
  # Returns `[stdout_and_stderr, exit_code]`
  def shell(*cmd, vars: nil)
    output, status, _ = _shell(*cmd, vars:)
    [output, status]
  end

  # like shell, but raises on non zero exit code. returns stdout_and_stderr otherwise
  def shell!(*cmd, vars: nil)
    output, status, cmd = _shell(*cmd, vars:)
    raise "#{cmd.inspect} failed #{status}\noutput: #{output}" if status != 0
    output
  end

  # one-liners
  def command?(cmd) = !!system("command", "-v", cmd, err: File::NULL)
  def md5(str) = Digest::MD5.hexdigest(str)
  def program_name = Pathname($PROGRAM_NAME).basename
  def sha256(str) = Digest::SHA256.hexdigest(str)

  #
  # banner/warning/fatal
  #

  GREEN = "\e[1;38;5;231;48;2;64;160;43m"
  YELLOW = "\e[1;38;5;231;48;2;251;100;11m"
  RED = "\e[1;38;5;231;48;2;210;15;57m"
  RESET = "\e[0m"

  def banner(str, color: GREEN)
    puts "#{color}[#{Time.new.strftime("%H:%M:%S")}] #{str.ljust(72)} #{RESET}"
  end

  def warning(str)
    banner(str, color: YELLOW)
  end

  def fatal(str)
    banner(str, color: RED)
    exit(1)
  end

  #
  # helpers
  #

  # Atomically replace a file by writing to a temporary path first.
  def atomic_write(path, &block)
    tmp = nil
    Pathname(path).tap do |path|
      path.dirname.mkpath
      tmp = Pathname("#{path}.tmp").tap { _1.unlink if _1.exist? }
      yield(tmp)
      tmp.rename(path)
    end
  ensure
    tmp.unlink if tmp&.exist?
  end

  # low-level helper for writing csv <= rows w/ headers
  def _csv_write0(csv, rows, headers: nil)
    headers ||= rows.first.to_h.keys
    csv << headers
    rows.each do |row|
      row = row.to_h
      csv << headers.map { row[_1] }
    end
  end

  # infer int/float from str
  def _infer_csv(str)
    case str
    when /\A-?\d+\z/ then return str.to_i
    when /\A-?\d+[.\d]+\z/ then return str.to_f
    end
    str
  end

  # shell helper
  def _shell(*cmd, vars: nil)
    begin
      cmd = cmd.first if cmd.one? && (cmd.first.is_a?(Array) || cmd.first.is_a?(String))
      if vars
        cmd = vars.reduce(cmd) do |memo, (k, v)|
          v = v.send(v.is_a?(Pathname) ? :escape : :to_s)
          memo.gsub("{{#{k}}}", v)
        end
      end
      cmd = Array(cmd).map { _1.is_a?(Pathname) ? _1.to_s : _1 }
      output, status = Open3.capture2e(*cmd)
      status = status.exitstatus
    rescue Errno::ENOENT => ex
      output, status = ex.message, 127
    end
    [output.strip, status, cmd]
  end
end

BTW - this could be cleaned up quite a bit if we could make some oft requested improvements to Pathname - ruby/pathname#64