grep-around/grep-around.rb

#!/usr/bin/env ruby
require 'optparse'

args = Array.new
rg_args = Array.new
add_to_rg_args = false

ARGV.each do |arg|
  if arg == "--"
    add_to_rg_args = true
  else
    if add_to_rg_args
      rg_args << arg
    else
      args << arg
    end
  end
end

options = {}
OptionParser.new do |opts|
  opts.banner = "Usage: grep-around.rb [options] -- [rg options]\n" \
    "All options after -- will be passed through to ripgrep"
  options[:stdin] = false
  options[:around] = 1

  opts.on("--stdin", "Read input from standard in") do |stdin|
    options[:stdin] = stdin
  end

  opts.on("--around [LINES]", "Lines around") do |around|
    options[:around] = around.to_i
  end

  opts.on("--up [LINES]", "Lines up") do |up|
    options[:up] = up.to_i
  end

  opts.on("--down [LINES]", "Lines down") do |down|
    options[:down] = down.to_i
  end
end.parse!(args)

%i[up down].each { |n| options[n] = options[:around] } if options[:around]

GREEN = "\u001b[32m"
RED = "\u001b[31m"
RESET = "\u001b[0m"

class GrepAround
  def initialize(options, rg_args)
    @options = options
    @rg_args = rg_args
    @hit = nil
    @current_line_number = nil
    @file_cache = Hash.new
  end

  def run
    results = run_ripgrep
    blocks = Array.new
    results.split("\n").each do |line|
      @hit = if line.match(/\A\d+:/)
               {
                 :file_name => :stdin,
                 :line_number => line.match(/\A\d+/)[0].to_i
               }
             else
               parse_metadata(line)
             end
      blocks << get_block
    end
    puts blocks.join("\n\n")
  end

  def run_ripgrep
    if @options[:stdin]
      contents = STDIN.read
      @file_cache[:stdin] = contents.split("\n")
      `echo '#{contents}' | rg -n #{@rg_args.join(' ')}`
    else
      `rg -n #{@rg_args.join(' ')}`
    end
  end

  def get_contents(key)
    if @file_cache.has_key?(key)
      @file_cache[key]
    else
      contents = File.readlines(@hit[:file_name])
      @file_cache[@hit[:file_name]] = contents
      contents
    end
  end

  def get_lines(contents, line_number)
    starting_line_number = line_number - @options[:up] - 1
    ending_line_number = line_number + @options[:down] - 1
    contents[starting_line_number..ending_line_number]
  end

  def parse_metadata(source)
    metadata = source.split(': ').first.split(':')
    file = metadata[0]
    line = metadata[1]
    {
      :file_name => file,
      :line_number => line.to_i,
    }
  end

  def get_block
    contents = get_contents(@hit[:file_name])
    lines = get_lines(contents, @hit[:line_number])

    block = Array.new
    block << block_header
    @current_line_number = @hit[:line_number] - @options[:up]
    lines.each do |line|
      block << build_output(line, @current_line_number)
      @current_line_number += 1
    end

    block.join("\n")
  end

  def block_header
    "#{RED}#{@hit[:file_name]}:#{RESET}"
  end

  def build_output(line, line_number)
    output = String.new
    output << line_number.to_s
    output << ' '
    output << GREEN if is_target_line
    output << line.gsub("\n", "")
    output << RESET if is_target_line

    output
  end

  def is_target_line
    @current_line_number == @hit[:line_number]
  end
end

GrepAround.new(options, rg_args).run