From 529a70e01b1f6d784ef0192850b9b5765b0787b2 Mon Sep 17 00:00:00 2001
From: Kimiyuki Onaka <kimiyuki95@gmail.com>
Date: Tue, 19 Mar 2019 21:05:51 +0900
Subject: [PATCH] #380: fix AtCoderProblem.get_input_format()

---
 onlinejudge/service/atcoder.py | 22 +++++++------
 tests/service_atcoder.py       | 57 ++++++++++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/onlinejudge/service/atcoder.py b/onlinejudge/service/atcoder.py
index f8a11f50..dce0f281 100644
--- a/onlinejudge/service/atcoder.py
+++ b/onlinejudge/service/atcoder.py
@@ -496,15 +496,19 @@ def get_input_format(self, session: Optional[requests.Session] = None) -> Option
 
         # 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'):
-                        return tag.decode_contents(formatter='html')
+        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/tests/service_atcoder.py b/tests/service_atcoder.py
index f5768254..15f3d5f2 100644
--- a/tests/service_atcoder.py
+++ b/tests/service_atcoder.py
@@ -177,6 +177,26 @@ def test_get_source_code(self):
 
 class AtCoderProblemGetInputFormatTest(unittest.TestCase):
     def test_normal(self):
+        """
+        .. 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')
@@ -186,8 +206,41 @@ def test_normal(self):
         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')
 
 
 if __name__ == '__main__':