PKUサポートスクリプト

学内ICPC練習会でPKU Online Judgeを使って練習をしているのだが、サンプル入力と出力をこぴぺするのが面倒に思えてきた。こういう単純作業はスクリプトに任せるのが定石だと思い、最近使い始めたpythonでさらっと書いてみた。

#!/usr/bin/python
import urllib
import os
import re
import sys
import subprocess
from optparse import OptionParser

options = None

def generateInputFileName(problemId, index):
    return problemId + '.' + str(index) + '.in.txt'

def generateOutputFileName(problemId, index):
    return problemId + '.' + str(index) + '.out.txt'

def format_pre(s):
    s = s.replace('<br />', '\n')
    s = s.replace('&lt;', '<')
    s = s.replace('&gt;', '>')
    s = s.replace('\r', '')
    if not s.endswith('\n'):
        s += '\n'
    while s.startswith('\n'):
        s = s[1:]
    return s

def download(id):
    global options
    url = 'http://acm.pku.edu.cn/JudgeOnline/problem?id=' + id
    html = None
    if options.titech_pubnet:
        html = urllib.urlopen(url, proxies={'http': 'http://proxy.noc.titech.ac.jp:3128'}).read()
    else:
        html = urllib.urlopen(url).read()
    p = re.compile('<pre class="sio">(.+?)</pre>', re.M | re.S)
    result = p.findall(html)
    n = len(result) / 2;
    for i in range(n):
        input_filepath = generateInputFileName(id, i)
        output_filepath = generateOutputFileName(id, i)
        open(input_filepath, 'w').write(format_pre(result[i * 2]))
        open(output_filepath, 'w').write(format_pre(result[i * 2 + 1]))
    return True

def compile(id):
    return subprocess.call(['g++', '-O2', '-o', 'a.exe', id + '.cpp']) == 0

def check(id):
    print 'compiling...'
    if not compile(id):
        print 'CompileError'
        exit(-1)
    
    if not os.path.exists(id + '.0.in.txt') or not os.path.exists(id + '.0.out.txt'):
        print 'downloading...'
        download(id)
    
    for i in range(100):
        input_filepath = generateInputFileName(id, i)
        output_filepath = generateOutputFileName(id, i)
    
        if os.path.exists(input_filepath):
            print 'Case #' + str(i) + ':'
    
            if os.path.exists(output_filepath):
                p = subprocess.Popen(['./a.exe'], stdin=open(input_filepath, 'r'), stdout=open('out.txt', 'w'))
            else:
                p = subprocess.Popen(['./a.exe'], stdin=open(input_filepath, 'r'))
    
            if p.wait() != 0:
                print 'RuntimeError?'
                exit(-1)
    
            if os.path.exists(output_filepath):
                if subprocess.call(['diff', output_filepath, 'out.txt', '-wy', '-W', '79']) != 0:
                    print 'WrongAnswer'
        else:
            break
    
    print 'OK'

def add_test_case_template(id):
    for file_index in range(100):
        input_filepath = generateInputFileName(id, file_index)
        output_filepath = generateOutputFileName(id, file_index)
        if os.path.isfile(input_filepath):
            continue
        open(input_filepath, 'w').close()
        open(output_filepath, 'w').close()
        print "Test case template file " + str(file_index) + " is created."
        return

def main():
    global options
    parser = OptionParser()
    parser.add_option("-t", "--titech-pubnet",
                      action="store_true", dest="titech_pubnet", default=False,
                      help="Use titech pubnet proxy",)
    parser.add_option("-a", "--add-test-case-template",
                      action="store_true", dest="add_test_case", default=False,
                      help="Generate a test case template file")
    parser.add_option("-c", "--check",
                      action="store_true", dest="check", default=True,
                      help="Build and check a solution")
    (options, args) = parser.parse_args()
    
    if len(args) == 0:
        print "Select problem id."
        parser.print_help()
        return

    id = args[0]

    if options.add_test_case:
        add_test_case_template(id)
        return
    
    if options.check:
        check(id)
        return
    
if __name__ == '__main__':
    main()

使い方は、スクリプトと同じ場所に「(問題番号).cpp」というソースを置いて、第一引数に問題番号を渡してスクリプトを実行するだけ。これでソースのコンパイル、サンプル入出力のダウンロード、解答チェックをやってくれる。Javaへの対応、TLEへの対応などがまだだが、まぁいいや。

追記(2010/09/04)

リファクタリングしました

追記(2010/09/24)

オプションでtitech-pubnet経由での接続と、入出力用空ファイルの生成に対応しました