atcoder_util/
sample_cases.rs

1extern crate reqwest;
2extern crate scraper;
3
4use scraper::{Html, Selector};
5use std::fs;
6use std::path::PathBuf;
7
8/// Struct to contain input/output examples of a problem.
9pub struct SampleCases {
10    pub input: Vec<String>,
11    pub output: Vec<String>,
12}
13
14impl SampleCases {
15    /// Construct a new `SampleCases`.
16    pub fn new(html: &Html) -> SampleCases {
17        let mut sc = SampleCases {
18            input: Vec::new(),
19            output: Vec::new(),
20        };
21        let io_examples = SampleCases::parse_io_examples(html);
22        sc.extract_io_example(io_examples);
23        sc
24    }
25
26    /// Construct a new `SampleCases` from input/output example files.
27    /// This method is called in `atcoder-util test`.
28    pub fn new_from_files(problem_id: &str) -> Self {
29        let mut sc = SampleCases {
30            input: Vec::new(),
31            output: Vec::new(),
32        };
33
34        let mut dir_io_examples = PathBuf::new();
35        dir_io_examples.push(format!("io_examples/{}", problem_id));
36        for dir_res in fs::read_dir(dir_io_examples).unwrap() {
37            let dir_name = dir_res.unwrap().path();
38            for file_res in fs::read_dir(&dir_name).unwrap() {
39                let file = file_res.unwrap();
40                let file_content = fs::read_to_string(file.path()).unwrap();
41                if dir_name.ends_with(&format!("{}_input", problem_id)) {
42                    sc.input.push(file_content);
43                } else {
44                    sc.output.push(file_content);
45                }
46            }
47        }
48        sc
49    }
50
51    /// Find elements of i/o examples in html and get them.
52    fn parse_io_examples(html: &Html) -> Vec<String> {
53        let selector_lang_ja = Selector::parse("span.lang-ja").unwrap();
54        let selector_io_example = Selector::parse("pre").unwrap();
55
56        let html_lang_ja = html.select(&selector_lang_ja).nth(0).unwrap().html();
57        let html_io_example = Html::parse_fragment(&html_lang_ja);
58
59        let io_examples: Vec<String> = html_io_example
60            .select(&selector_io_example)
61            .filter(|example| example.children().count() == 1)
62            .map(|example| example.text().collect::<Vec<&str>>().join(""))
63            .collect();
64
65        io_examples
66    }
67
68    /// Push i/o examples to vector in a struct itself.
69    fn extract_io_example(&mut self, io_examples: Vec<String>) {
70        for (i, io_example) in io_examples.iter().enumerate() {
71            // IO example of even index is input.
72            if i % 2 == 0 {
73                self.input.push(io_example.to_string());
74            } else {
75                self.output.push(io_example.to_string());
76            }
77        }
78    }
79}