[usaco] Add a support for several sample tests

niyaznigmatullin · Aug 15, 2020 · f030537 · f030537
1 parent f8f6d3c
commit f030537
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/polygon_uploader/__init__.py b/polygon_uploader/__init__.py
@@ -1,3 +1,3 @@
 
 
-__version__ = "0.2.6"
+__version__ = "0.2.7"
diff --git a/polygon_uploader/usaco/usaco.py b/polygon_uploader/usaco/usaco.py
@@ -92,10 +92,11 @@ def extract_part(class_name):
             for x in statement.find_all_next('h4'):
                 x.extract()
 
+            sample_count = len(statement.find_all_next('pre', attrs={'class', 'in'}))
             ps = statement.find_all_next('p')
             note = None
             for e in ps[::-1]:
-                if len(e.find_all_next('pre', attrs={'class', 'in'})) > 0:
+                if len(e.find_all_next('pre', attrs={'class', 'in'})) == sample_count:
                     e.extract()
                     note = e
                     break
@@ -118,8 +119,9 @@ def extract_part(class_name):
                                           scoring=scoring,
                                           notes=note)
             prob.save_statement(lang=lang_polygon, problem_statement=polygon_statement)
+        return sample_count
 
-    def download_tests(dir, tests_dir):
+    def download_tests(dir, tests_dir, sample_count):
         tests_archive = os.path.join(dir, "tests.zip")
         download_file_to(testdata_href, tests_archive)
         print(tests_archive, "downloaded")
@@ -134,8 +136,8 @@ def file_to_test(name):
             return FileTest(os.path.join(tests_dir, name), 'usacoimport: filename = %s' % name)
 
         groups = [
-            Group(0, [file_to_test('1.in')], GroupScoring.SUM),
-            Group(100, [file_to_test('%d.in' % x) for x in range(2, cnt + 1)], GroupScoring.SUM),
+            Group(0, [file_to_test('%d.in' % x) for x in range(1, sample_count + 1)], GroupScoring.SUM),
+            Group(100, [file_to_test('%d.in' % x) for x in range(sample_count + 1, cnt + 1)], GroupScoring.SUM),
         ]
 
         upload_groups(prob, groups)
@@ -183,11 +185,12 @@ def download_solutions():
     print("problem.enableGroups")
     prob.enable_groups('tests', True)
 
+    sample_count = download_statement()
+
     tests_dir = os.path.join(dir, "tests")
     os.mkdir(tests_dir)
-    download_tests(dir, tests_dir)
+    download_tests(dir, tests_dir, sample_count)
 
-    download_statement()
     download_solutions()
 
     print("problem.setChecker std::wcmp.cpp")