【競プロ】サンプルテストケースを取得してくるシェルスクリプトを書いた

前の記事でinファイルにサンプルテストケースの入力例と出力例を自動で書き込めるようにしたいと書きましたが、できるようになりました。 現時点ではAtcoder、AOJ（Arenaを除く）、yukicoderに対応しています。今回のも前回と同様githubに上げてあります。

https://github.com/hpano/CompetitiveProgramming/blob/master/add.sh

動作はこんな感じ。例としてAtcoderのABC173のA問題とB問題のサンプルテストケースを取得してきています。

コンテストページはこちら

atcoder.jp

さて順番にコードを見ていきましょう。

まず、-rオプションがある場合はすでにあるinファイルを削除するので、フラグを立てます。そして作りたいファイル名を順に配列に突っ込んでいきます。

declare -a argv=()
declare rflag=false
while test "$1" != ""; do
  case $1 in
    -*)
      if [[ "$1" =~ 'r' ]] || [[ "$1" =~ '-remove' ]]; then
        rflag=true
      fi
      shift
      ;;
    *)
      argv=("${argv[@]}" "${1%.cpp}")
      shift
      ;;
  esac
done

そして各ファイル名のC++ファイルを、テンプレートC++ファイルをコピーして作成します。またinファイル用のデータを取得するために、パスのディレクトリ名から競プロサイトを判断して後の処理につなげます。

この先はどのサイトでも大筋は同じなので、例としてAtcoderの場合を見ていきましょう。

for name in ${argv[@]}; do
  if [ -f ${name}.cpp ]; then
    echo "[add] ${name}.cpp already exists"
    continue
  fi

  cp $(dirname $0)/template.cpp ${name}.cpp && \
  echo "[add] Added ${name}.cpp"
done

declare path=( $(pwd | sed -e "s/.*competitive_programming\///g" | tr "/" " ") )
case ${path[0]} in
  "atcoder")
    get_atcoder
    ;;
  "AOJ")
    get_aoj
    ;;
  "yukicoder")
    get_yukicoder
    ;;
  *)
    echo "[add] ${path[0]} is unsupported"
    exit 1
    ;;
esac

まず実行ディレクトリまでのパスからコンテスト名を判断していき、URLに足していきます。例えばABC173の場合は次のような感じです。

"https://atcoder.jp/contests/" + "abc" + "173" + "/tasks/" + "abc" + "173_" + "問題ID[a-f]"

Atcoderの場合はABC、ARC、AGCは概ねこの規則通りなのでいいのですが、企業コンなどそれ以外の場合はちょっと異なるので、上記と異なる処理が必要です。

まずは上記と同じように処理していき、curlを使って一度HTTPリクエストを投げます。これでデータが取得できればそれでOKなのですが、できなかった場合、URLの最後に付いてるbeginnerを除去してもう一度curlします。URL構築の方法上、コンテスト名と同じものが"/tasks/"の後に追加されるのですが、コンテスト名にbeginnerが付いている場合でもURLの"/tasks/"の後には付いていないので（例えば下のコンテストなど）。

atcoder.jp

それでも404が返ってくる場合、readコマンドでURLの最後の部分を直接入力させます。

function get_atcoder() {
  local url_pre="https://atcoder.jp/contests/"
  # URLの前処理 企業コンなどotherの場合は特別
  if [ ${path[1]} = "other" ]; then
    url_pre+=$(echo ${path[2]} | tr "[A-Z]" "[a-z]")
    url_pre+="/tasks/"
    url_pre+=$(echo ${path[2]}_ | tr "[A-Z-]" "[a-z_]")
    local tmp_status=$(curl -Ls ${url_pre}a -o /dev/null -w '%{http_code}\n')
    if [ ${tmp_status} -eq 404 ]; then
      if [[ ${url_pre} =~ "beginner" ]]; then
        url_pre=${url_pre%%_beginner*}
        url_pre+="_"
      fi
      tmp_status=$(curl -Ls ${url_pre}a -o /dev/null -w '%{http_code}\n')
    fi
    if [ ${tmp_status} -eq 404 ]; then
      local url_in
      url_pre="${url_pre%\/*}/"
      read -p "URL? > ${url_pre}" url_in
      if [ "${url_in: -1}" != "_" ]; then
        if [[ "${url_in: -2}" =~ _[a-z] ]]; then
          url_in=${url_in%_*}
        fi
        url_in+="_"
      fi
      url_pre+=${url_in}
    fi
  else
    url_pre+=$(echo ${path[1]} | tr "[A-Z]" "[a-z]")
    url_pre+=$(echo "${path[2]}/tasks/")
    url_pre+=$(echo ${path[1]} | tr "[A-Z]" "[a-z]")
    url_pre+=$(echo "${path[2]}_")
  fi

ここから入力された各ファイル名について処理していきます。

まず、すでにinファイルが存在する場合、-rオプションが入力されていればそのinファイルを削除し、入力されていなければcontinueして次のファイルに移ります。

  # 各問題に対してデータ取得とinファイル作成
  for name in ${argv[@]}; do
    if [ -f ${name}.in ]; then
      if "${rflag}"; then
        rm ${name}.in
      else
        echo "[add] ${name}.in already exists"
        continue
      fi
    fi

続いてURLの最終決定です。昔のABCとかだと問題IDが[a-f]ではなく[1-6]になっていたりするので、それに対応するためにurl2として用意します。urlを試して404の場合url2を試すといった感じですね。それでもステータスが200でなかった場合、最終的にここでURLを全て入力してもらいます。もしくは、"n"を入力するとその問題を飛ばすことができます。

    local url=${url_pre}
    url+=${name}
    local url2=${url_pre}
    url2+=$(echo ${name} | tr "[a-f]" "[1-6]")
    local url2flag=false
    local continue_flag=false
    local status=$(curl -Ls ${url} -o /dev/null -w '%{http_code}\n')
    if [ ${status} -eq 404 ]; then
      url2flag=true
      sleep 0.5
      status=$(curl -Ls ${url2} -o /dev/null -w '%{http_code}\n')
    fi
    while [ ${status} -ne 200 ]; do
      echo "[add] Cannot recieve contents from ${url} (code: ${status})"
      read -p "Input correct URL or 'n' to give up getting: " url
      status=$(curl -Ls ${url} -o /dev/null -w '%{http_code}\n')
      url2flag=false
      if [[ ${url} == "n" ]]; then
        continue_flag=true
        break
      fi
    done
    if "${continue_flag}"; then
      continue
    fi
    if "${url2flag}"; then
      url=${url2}
    fi

そして得られたHTMLから必要な範囲だけを抽出します。どの部分が必要かは各競プロサイトのソースを見てみましょう。

    local data=$(curl -Lso- ${url})
    data=${data#*<span class=\"lang-ja\">}
    data=${data%%</span>*}
    local inflag=false
    local outflag=false
    touch ${name}.in

    echo "[add] Get from ${url}"

最後に残ったHTMLを1行ずつ見ていき、入力例or出力例の部分かを判断して、各テストケース毎に次のようにinファイルに書き込んで終わりです。

<in>
入力例1
</in>
<out>
出力例1
</out>

    while read line; do
      line=`echo ${line} | tr -d "\r"`
      if "${inflag}"; then
        if [[ ${line} =~ "</pre>" ]]; then
          inflag=false
          echo "</in>" >> ${name}.in
        else
          echo ${line} >> ${name}.in
        fi
      elif "${outflag}"; then
        if [[ ${line} =~ "</pre>" ]]; then
          outflag=false
          echo "</out>" >> ${name}.in
        else
          echo ${line} >> ${name}.in
        fi
      elif [[ ${line} =~ "<h3>入力例" ]]; then
        echo "<in>" >> ${name}.in
        inflag=true
        echo ${line#*<pre>} >> ${name}.in
      elif [[ ${line} =~ "<h3>出力例" ]]; then
        echo "<out>" >> ${name}.in
        outflag=true
        echo ${line#*<pre>} >> ${name}.in
      fi
    done << END
    ${data}
END

    echo "[add] Added ${name}.in"
  done
}

サンプルテストケースをコピペする手間が省けてかなり便利になりました。