Merge pull request #211 from kmyk/issue/210

#210
online-judge-tools · Dec 27, 2018 · e53ae05 · e53ae05
2 parents 3fdce0c + 3482fda
commit e53ae05
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/onlinejudge/atcoder.py b/onlinejudge/atcoder.py
@@ -122,21 +122,32 @@ def _get_tag_lang(self, tag):
                 if cls.startswith('lang-'):
                     return cls
 
-    def _find_sample_tags(self, soup):
-        result = []
+    def _find_sample_tags(self, soup) -> Generator[Tuple[bs4.Tag, bs4.Tag], None, None]:
         for pre in soup.find_all('pre'):
             log.debug('pre tag: %s', str(pre))
             if not pre.string:
                 continue
             prv = utils.previous_sibling_tag(pre)
-            if prv and prv.name == 'h3' and prv.string:  # AtCoder's javascript recognizes `h3+pre' as a sample input/output
-                result += [( pre, prv )]
+
+            # the first format: h3+pre
+            if prv and prv.name == 'h3' and prv.string:
+                yield ( pre, prv )
+
             else:
-                if pre.parent and pre.parent.name == 'section':  # AtCoder's javascript sometimes fails. e.g. http://abc001.contest.atcoder.jp/tasks/abc001_1
+                # ignore tags which are not samples
+                # example: https://atcoder.jp/contests/abc003/tasks/abc003_4
+                while prv is not None:
+                    if prv.name == 'pre':
+                        break
+                    prv = utils.previous_sibling_tag(prv)
+                if prv is not None:
+                    continue
+
+                # the second format: h3+section pre
+                if pre.parent and pre.parent.name == 'section':
                     prv = pre.parent and utils.previous_sibling_tag(pre.parent)
                     if prv and prv.name == 'h3' and prv.string:
-                        result += [( pre, prv )]
-        return result
+                        yield ( pre, prv )
 
     def get_url(self) -> str:
         return 'http://{}.contest.atcoder.jp/tasks/{}'.format(self.contest_id, self.problem_id)

diff --git a/tests/download_atcoder.py b/tests/download_atcoder.py
@@ -94,3 +94,10 @@ def test_call_download_atcoder_abc114_c(self):
                 {"input": "3600\n", "output": "13\n"},
                 {"input": "999999999\n", "output": "26484\n"},
             ], type='json')
+    def test_call_download_atcoder_abc114_c(self):
+        self.snippet_call_download('https://atcoder.jp/contests/abc003/tasks/abc003_4', [
+                {"input": "3 2\n2 2\n2 2\n", "output": "12\n"},
+                {"input": "4 5\n3 1\n3 0\n", "output": "10\n"},
+                {"input": "23 18\n15 13\n100 95\n", "output": "364527243\n"},
+                {"input": "30 30\n24 22\n145 132\n", "output": "976668549\n"},
+            ], type='json')