LoginSignup
23
11

More than 3 years have passed since last update.

RustAdvent Calendar 2019Day 19
@togatoga(Hitoshi Togasaki)

RustでAtCoderのサンプルケースを自動取得する

Last updated at Posted at 2019-12-19

はじめに

この記事はRust Advent Calendar 2019 19日目の記事です。
AtCoderとは国内最大のプログラミングコンテストサイトです。競技プログラミング・プログラミングコンテストについてはdrkenさんのこちらの記事を参照してください。

AtCoderではコードをジャッジサーバーに提出する前に、手元で動作確認を行うための2~3個のサンプルケースが問題のページに記載されています。
AtCoder Beginner Contest 147 A - Blackjackより
Screen Shot 2019-12-20 at 1.38.10 AM.png

プログラムを修正するたびに問題のページからサンプルケースをコピペするのは手間でかつ、コピペミスの可能性もあるためサンプルケースは自動に取得して保存しておきたいです。
そこでAtCoderの問題ページをスクレイピングしてサンプルケースを保存するコマンドラインツールをRustで作成しました。
今回は作成したツールの説明と実装について軽く説明したいと思います。

atcoder-sample-downloader

今回、作成したツールはatcoder-sample-downloaderです。使い方などはGitHubのREADMEに書きました。downloadとloginの2つのサブコマンドが存在します。

download

問題のURLを引数に問題ページをスクレイピングしてサンプルケースを取得し、現在のディレクトリにサンプルケースのファイルを保存します。


% atcoder-sample-downloder download https://atcoder.jp/contests/abc147/tasks/abc147_a
====== Download Result ======
=== Sample Test Case 1 ===
Input:
5 7 9

Output:
win

=== Sample Test Case 2 ===
Input:
13 7 2

Output:
bust

=============================
% bat sample_*
───────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: sample_input_1.txt
───────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ 5 7 9
   2   │ 11 17
───────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────
───────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: sample_input_2.txt
───────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ 13 7 2
   2   │ 2 1
───────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────
───────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: sample_output_1.txt
───────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ win
───────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────
───────┬──────────────────────────────────────────────────────────────────────────────────────────────────────────────
       │ File: sample_output_2.txt
───────┼──────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1   │ bust
───────┴──────────────────────────────────────────────────────────────────────────────────────────────────────────────
実装編

使った主なライブラリは以下です。
1. コマンドラインパーサー
- clap
2. スクレイピング
- scraper
3. HTTPリクエスト
- reqwest

今回は2つのAtCoderとAtCoderParserというstructを作成しました。

struct AtCoder {
    client: reqwest::Client,
    //for request
    cookie_headers: HeaderMap,
    //from response
    html: Option<String>,
}

struct AtCoderParser {
    document: scraper::Html,
}

AtCoderParserはhtmlからサンプルケースやTokenの情報を取得します。

impl AtCoderParser {
    fn new(html: &str) -> AtCoderParser {
        AtCoderParser {
            document: scraper::Html::parse_document(html),
        }
    }
    fn sample_cases(&self) -> Option<Vec<(String, String)>> {
        let task_statement_selector =
            scraper::Selector::parse(r#"div[id="task-statement"]"#).unwrap();
        let pre_selector = scraper::Selector::parse("pre").unwrap();
        let h3_selector = scraper::Selector::parse("h3").unwrap();
        let input_h3_text = vec!["入力例", "Sample Input"];
        let output_h3_text = vec!["出力例", "Sample Output"];

        let mut input_cases = vec![];
        let mut output_cases = vec![];
        if let Some(task_statement) = self.document.select(&task_statement_selector).next() {
            for pre in task_statement.select(&pre_selector) {
                if let Some(pre_parent) = pre.parent_element() {
                    if let Some(h3) = pre_parent.select(&h3_selector).next() {
                        let h3_text = h3.text().collect::<String>();
                        let input = input_h3_text.iter().any(|&x| h3_text.contains(x));
                        let output = output_h3_text.iter().any(|&x| h3_text.contains(x));
                        let text = pre.text().collect::<String>();
                        if input {
                            input_cases.push(text);
                        } else if output {
                            output_cases.push(text);
                        }
                    }
                }
            }
        } else {
            return None;
        }
        // make cases unique to remove extra duplicated language cases
        let input_cases: Vec<String> = input_cases.into_iter().unique().collect();
        let output_cases: Vec<String> = output_cases.into_iter().unique().collect();
        let sample_test_cases: Vec<(String, String)> = input_cases
            .into_iter()
            .zip(output_cases)
            .map(|(input, output)| (input, output))
            .collect();
        Some(sample_test_cases)
    }

    fn csrf_token(&self) -> Option<String> {
        let selector = scraper::Selector::parse(r#"input[name="csrf_token"]"#).unwrap();
        if let Some(element) = self.document.select(&selector).next() {
            if let Some(token) = element.value().attr("value") {
                return Some(token.to_string());
            }
        }
        None
    }
}

AtCoderは実際にAtCoderのサイトにHTTPリクエストを行い、得られたリスポンスのCookieとhtmlを保存しておきます。保存したCookieの用途はloginで説明します。 


    pub async fn download(&mut self, url: &str) -> Result<(), failure::Error> {
        let url = url::Url::parse(url)?;
        if let Ok(cookie_headers) = AtCoder::local_cookie_headers() {
            self.cookie_headers = cookie_headers;
        }
        let resp = self.call_get_request(url.as_str()).await?;
        self.parse_response(resp).await?;

        let parser = AtCoderParser::new(&self.html.as_ref().unwrap());
        let sample_test_cases = parser.sample_cases();

        println!("====== Download Result ======");
        if let Some(samples) = sample_test_cases {
            AtCoder::create_sample_test_files(&samples)?;
            for (idx, (input, output)) in samples.iter().enumerate() {
                println!("=== Sample Test Case {} ===", idx + 1);
                println!("Input:\n{}\nOutput:\n{}", input, output);
            }
        }
        println!("=============================");
        Ok(())
    }

login

loginコマンドはターミナル上でユーザー名とパスワードを入力してAtCoderにログインし、取得したログイン情報(Cookie)をローカルに保存します。
問題のHTMLを取得する際に保存したCookie情報を使うことでログイン状態を保ったままHTMLを取得できます。開催中のコンテスト問題は現在参加しているユーザーのみがアクセス可能です。Cookie情報を使うことで開催中のコンテスト問題のサンプルケースも取得することができます。(ログインしているユーザーが参加していること)
ログインが伴うウェブサイトのスクレイピングはHow To Scrape A Website That Requires Login With Golang?が大変参考になりました。

% atcoder-sample-downloder login
Please input Your username and password
Username >
Password >
SAVED YOUR COOKIE IN /Users/<USER NAME>/.atcoder-sample-downloader/cookie.jar
実装編

AtCoderのログインページにログインに必要な情報をPOSTする必要があります。
ブラウザーのDeveloperツールを使うことで、実際にどんな情報を送信しているかを確認することができます。
画像はChromeのDeveloperツールです。Screen Shot 2019-12-20 at 1.13.45 AM.png
Developerツールからユーザー名とパスワードとcsrf_tokenとcsrf_token取得時のCookieも併せてPOSTする必要があることがわかりました。
以下が実際のコードです。


    pub async fn login(&mut self, url: &str) -> Result<(), failure::Error> {
        let url = url::Url::parse(url)?;
        let resp = self.call_get_request(url.as_str()).await?;
        self.parse_response(resp).await?;
        let parser = AtCoderParser::new(self.html.as_ref().unwrap());
        //necessary information and parameters to login AtCoder
        let csrf_token = parser.csrf_token().unwrap();
        let (username, password) = AtCoder::username_and_password();
        let params = {
            let mut params = std::collections::HashMap::new();
            params.insert("username", username);
            params.insert("password", password);
            params.insert("csrf_token", csrf_token);
            params
        };
        //make a post request and try to login
        let resp = self.call_post_request(url.as_str(), &params).await?;
        //save your cookie in your local
        AtCoder::save_cookie_in_local(&resp)?;
        Ok(())
    }

最後に

実は元々はGoでサンプルケースの自動取得や実行などのプロコンツールを書いていました。今年から本格的にRustを勉強し始めたのでRustでツールを書き直すことも検討していたのですが、本当にできるのかと思い実験も兼ねて挑戦してみました。来年からはRustでどんどんツールを作っていきたいと思います！！

23
11
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
23
11