sample_getter.py

# -*- coding: utf-8 -*-
import sys
import urllib2
from bs4 import BeautifulSoup

def sample_getter(path, problemId):
    directory = path + "/" + str(problemId)
    url = "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=" + str(problemId)
    html = urllib2.urlopen(url)
    soup = BeautifulSoup(html)
    
    input_element = [e for e in soup.find_all("h2") if e.string == "Sample Input"][0]
    output_element = \
        [e for e in soup.find_all("h2") if e.string == "Output for the Sample Input"][0]
    
    f1 = open(directory + "/" + str(problemId) + "_sample_input.txt", 'w')
    f1.write(input_element.next_element.next_element.next_element.string)
    f1.close()
    
    f2 = open(directory + "/" + str(problemId) + "_sample_output.txt", 'w')
    f2.write(output_element.next_element.next_element.next_element.string)
    f2.close()

if __name__ == '__main__':
    problemId = sys.argv[1]
    sample_getter("../AOJ", problemId)