Repository: kmyk/online-judge-tools (python)

Issue:
Fails to fetch the input formats of old Atcoder problems
昔のAtCoderの問題、例えば https://atcoder.jp/contests/arc002/tasks/arc002_3 のような問題の入力フォーマットのフェッチに失敗します。(''が返ってくる)
これはhtmlの構造が昔と今で異なるからだと思いますが、サポートするのは良いと思いますか?



Pull request description:
#380: fix AtCoderProblem.get_input_format()
close #380

なんだか複数になってしまったが、やったのは以下です。

-   #380 で議論した形で仕様を固定
-   #380 のバグを修正
-   その他のバグもついでに修正

Patch:
diff --git a/onlinejudge/service/atcoder.py b/onlinejudge/service/atcoder.py
index d0fde48..dce0f28 100644
--- a/onlinejudge/service/atcoder.py
+++ b/onlinejudge/service/atcoder.py
@@ -481,7 +481,7 @@ class AtCoderProblem(onlinejudge.type.Problem):
 
         return None
 
-    def get_input_format(self, session: Optional[requests.Session] = None) -> str:
+    def get_input_format(self, session: Optional[requests.Session] = None) -> Optional[str]:
         """
         :raises Exception: if no such problem exists
         """
@@ -496,19 +496,20 @@ class AtCoderProblem(onlinejudge.type.Problem):
 
         # parse
         soup = bs4.BeautifulSoup(resp.content.decode(resp.encoding), utils.html_parser)
-        for h3 in soup.find_all('h3'):
-            if h3.string in ('入力', 'Input'):
-                tag = h3
-                for _ in range(3):
-                    tag = utils.next_sibling_tag(tag)
-                    if tag is None:
-                        break
-                    if tag.name in ('pre', 'blockquote'):
-                        s = ''
-                        for it in tag:
-                            s += it.string or it  # AtCoder uses <var>...</var> for math symbols
-                        return s
-        return ''
+        for h3 in soup.find_all('h3', text=re.compile(r'^(入力|Input)$')):
+            if h3.parent.name == 'section':
+                section = h3.parent
+            else:
+                section = h3.find_next_sibling('section')
+            if section is None:
+                section = soup.find(class_='io-style')
+            if section is None:
+                log.warning('<section> tag not found. something wrong')
+                return None
+            pre = section.find('pre')
+            if pre is not None:
+                return pre.decode_contents(formatter=None)
+        return None
 
     def get_available_languages(self, session: Optional[requests.Session] = None) -> List[Language]:
         """
diff --git a/onlinejudge/service/yukicoder.py b/onlinejudge/service/yukicoder.py
index 49b8f53..3472f90 100644
--- a/onlinejudge/service/yukicoder.py
+++ b/onlinejudge/service/yukicoder.py
@@ -420,7 +420,7 @@ class YukicoderProblem(onlinejudge.type.Problem):
         soup = bs4.BeautifulSoup(resp.content.decode(resp.encoding), utils.html_parser)
         for h4 in soup.find_all('h4'):
             if h4.string == '入力':
-                return h4.parent.find('pre').string
+                return h4.parent.find('pre').decode_contents(formatter=None)
         return None
 
 
diff --git a/onlinejudge/type.py b/onlinejudge/type.py
index 6cc5cb3..d90257c 100644
--- a/onlinejudge/type.py
+++ b/onlinejudge/type.py
@@ -111,6 +111,10 @@ class Problem(object):
         raise NotImplementedError
 
     def get_input_format(self, session: Optional[requests.Session] = None) -> Optional[str]:
+        """
+        :return: the HTML in the `<pre>` tag as :py:class:`str`
+        """
+
         raise NotImplementedError
 
     @classmethod


Test patch:
diff --git a/tests/service_atcoder.py b/tests/service_atcoder.py
index 29b45ef..3ab68f2 100644
--- a/tests/service_atcoder.py
+++ b/tests/service_atcoder.py
@@ -177,17 +177,76 @@ class AtCoderSubmissionTest(unittest.TestCase):
 
 class AtCoderProblemGetInputFormatTest(unittest.TestCase):
     def test_normal(self):
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc001/tasks/agc001_d').get_input_format(), 'N M\r\nA_1 A_2 ... A_M\r\n')
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc002/tasks/agc002_d').get_input_format(), '\r\nN M\r\na_1 b_1\r\na_2 b_2\r\n:\r\na_M b_M\r\nQ\r\nx_1 y_1 z_1\r\nx_2 y_2 z_2\r\n:\r\nx_Q y_Q z_Q\r\n')
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc003/tasks/agc003_d').get_input_format(), 'N\r\ns_1\r\n:\r\ns_N\r\n')
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc004/tasks/agc004_d').get_input_format(), 'N K\r\na_1 a_2 ... a_N\r\n')
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc005/tasks/agc005_d').get_input_format(), 'N K\r\n')
-
-        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/arc083/tasks/arc083_a').get_input_format(), 'A B C D E F\r\n')
+        """
+        .. code-block:: html
+
+            <div class="io-style">
+                <div class="part">
+                    <section>
+                        <h3>入力</h3>
+                        <p>入力は以下の形式で標準入力から与えられる。</p>
+                        <pre>
+                            <var>N</var>
+                        </pre>
+                    </section>
+                </div>
+                <div class="part">
+                    ...
+                </div>
+                ...
+            </div>
+        """
+
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc001/tasks/agc001_d').get_input_format(), '<var>N</var> <var>M</var>\r\n<var>A_1</var> <var>A_2</var> <var>...</var> <var>A_M</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc002/tasks/agc002_d').get_input_format(), '\r\n<var>N</var> <var>M</var>\r\n<var>a_1</var> <var>b_1</var>\r\n<var>a_2</var> <var>b_2</var>\r\n<var>:</var>\r\n<var>a_M</var> <var>b_M</var>\r\n<var>Q</var>\r\n<var>x_1</var> <var>y_1</var> <var>z_1</var>\r\n<var>x_2</var> <var>y_2</var> <var>z_2</var>\r\n<var>:</var>\r\n<var>x_Q</var> <var>y_Q</var> <var>z_Q</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc003/tasks/agc003_d').get_input_format(), '<var>N</var>\r\n<var>s_1</var>\r\n:\r\n<var>s_N</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc004/tasks/agc004_d').get_input_format(), '<var>N</var> <var>K</var>\r\n<var>a_1</var> <var>a_2</var> <var>...</var> <var>a_N</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/agc005/tasks/agc005_d').get_input_format(), '<var>N</var> <var>K</var>\r\n')
+
+        self.assertEqual(AtCoderProblem.from_url('https://beta.atcoder.jp/contests/arc083/tasks/arc083_a').get_input_format(), '<var>A</var> <var>B</var> <var>C</var> <var>D</var> <var>E</var> <var>F</var>\r\n')
 
     def test_old_problem(self):
-        # https://github.com/kmyk/online-judge-tools/issues/380
-        pass
+        """
+        :note: https://github.com/kmyk/online-judge-tools/issues/380
+
+        .. code-block:: html
+
+            <h3>入力</h3>
+            <section>
+                入力は以下の形式で与えられる。
+                <pre>
+                    <var>N</var>
+                </pre>
+            </section>
+        """
+
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/arc001/tasks/arc001_1').get_input_format(), '\r\n<var>N</var>\r\n<var>c_1c_2c_3…c_N</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/arc002/tasks/arc002_3').get_input_format(), '\r\n<var>N</var>\r\n<var>c_{1}c_{2}...c_{N}</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/arc034/tasks/arc034_4').get_input_format(), '\r\n<var>A</var> <var>B</var> <var>C</var>\r\n<var>a_1</var> <var>a_2</var> .. <var>a_A</var>\r\n<var>b_1</var> <var>b_2</var> .. <var>b_B</var>\r\n')
+
+    def test_dwacon_problem(self):
+        """
+        :note: https://github.com/kmyk/online-judge-tools/issues/142
+
+        .. code-block:: html
+
+            <h3 id="入力">入力</h3>
+            <p>入力は以下の形式で標準入力から与えられる。</p>
+            <div class="io-style">
+                <pre>
+                    <var>N</var>
+                </pre>
+            </div>
+        """
+
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/dwacon2018-final/tasks/dwacon2018_final_a').get_input_format(), '\r\n<var>H</var> <var>M</var> <var>S</var>\r\n<var>C_1</var> <var>C_2</var>\r\n')
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/dwacon2018-final/tasks/dwacon2018_final_b').get_input_format(), '\r\n<var>N</var> <var>K</var>\r\n<var>v_1</var> <var>...</var> <var>v_N</var>\r\n')
+
+    def test_problem_without_input(self):
+        self.assertIsNone(AtCoderProblem.from_url('https://atcoder.jp/contests/tenka1-2013-quala/tasks/tenka1_2013_qualA_a').get_input_format())
+
+    def test_problem_without_input_format(self):
+        self.assertIsNone(AtCoderProblem.from_url('https://atcoder.jp/contests/joi2006ho/tasks/joi2006ho_a').get_input_format())
 
 
 if __name__ == '__main__':
diff --git a/tests/service_yukicoder.py b/tests/service_yukicoder.py
index 1c15a74..14a02a1 100644
--- a/tests/service_yukicoder.py
+++ b/tests/service_yukicoder.py
@@ -148,6 +148,9 @@ class YukicoderProblemGetInputFormatTest(unittest.TestCase):
         self.assertEqual(YukicoderProblem.from_url('https://yukicoder.me/problems/no/512').get_input_format(), '$X$ $Y$\n$N$\n$A_1$ $\\cdots$ $A_N$\n')
         self.assertEqual(YukicoderProblem.from_url('https://yukicoder.me/problems/no/777').get_input_format(), '$N$\n$A_1$ $B_1$ $C_1$\n$A_2$ $B_2$ $C_2$\n…\n$A_N$ $B_N$ $C_N$\n')
 
+    def test_problem_without_input(self):
+        self.assertIsNone(YukicoderProblem.from_url('https://yukicoder.me/problems/no/3003').get_input_format())
+
 
 if __name__ == '__main__':
     unittest.main()