Downloader for input-data of online-judge-problem (programming-contest)

input_gen.rb

#!/usr/bin/env ruby
# -*- mode: ruby -*-
require 'net/http'
require 'uri'
require 'nokogiri'

module AOJ
  extend self

  def get(id)
    id = "0" * (4 - id.to_s.length) + id.to_s
    url = URI.parse("http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=#{id}")
    html = Net::HTTP.get(url)
    doc = Nokogiri::HTML.parse(html, nil, nil)
  end

  def extract(doc)
    sample_inputs, sample_outputs = [], []
    doc.search("//h2").each do |h2|
      if h2.text =~ /^Sample Input ?(\d)?$/
        sample_inputs << [h2.text, $1, h2.next_element.text.strip.chomp]
      elsif h2.text =~ /^Output for the Sample Input ?(\d)?$/
        sample_outputs << [h2.text, $1, h2.next_element.text.strip.chomp]
      end
    end
    return sample_inputs, sample_outputs
  end

  def in_gen(id)
    extract(get(id))[0].each do |title, num, input|
      puts "[#{title}]"
      puts input
      `echo '#{input}' | cat > in#{num}`
    end
  end
end

module POJ


end

def parse(problem_id)
  if problem_id =~ /^([a-z]+)(\d+)$/
    case $1
    when "aoj"
      return :aoj, $2
    when "poj"
      return :poj, $2
    end
  end
  raise "unknown problem id #{problem_id}"
end


site, id = parse(ARGV[0])
case site
when :aoj
  AOJ.in_gen(id)
when :poj
  POJ.ins_gen(id)
end