Fix crash bug of envgen when getting 404 (#201)

kyuridenamida · May 16, 2020 · 3d04f49 · 3d04f49
1 parent 8b71b65
commit 3d04f49
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/atcodertools/client/atcoder.py b/atcodertools/client/atcoder.py
@@ -8,19 +8,23 @@
 import requests
 from bs4 import BeautifulSoup
 
+from atcodertools.client.models.contest import Contest
+from atcodertools.client.models.problem import Problem
+from atcodertools.client.models.problem_content import ProblemContent, InputFormatDetectionError, SampleDetectionError
 from atcodertools.client.models.submission import Submission
 from atcodertools.common.language import Language
 from atcodertools.common.logging import logger
 from atcodertools.fileutils.artifacts_cache import get_cache_file_path
-from atcodertools.client.models.contest import Contest
-from atcodertools.client.models.problem import Problem
-from atcodertools.client.models.problem_content import ProblemContent, InputFormatDetectionError, SampleDetectionError
 
 
 class LoginError(Exception):
     pass
 
 
+class PageNotFoundError(Exception):
+    pass
+
+
 default_cookie_path = get_cache_file_path('cookie.txt')
 
 
@@ -108,6 +112,8 @@ def login(self,
     def download_problem_list(self, contest: Contest) -> List[Problem]:
         resp = self._request(contest.get_problem_list_url())
         soup = BeautifulSoup(resp.text, "html.parser")
+        if resp.status_code == 404:
+            raise PageNotFoundError
         res = []
         for tag in soup.find('table').select('tr')[1::]:
             tag = tag.find("a")

diff --git a/atcodertools/tools/envgen.py b/atcodertools/tools/envgen.py
@@ -11,7 +11,7 @@
 
 from colorama import Fore
 
-from atcodertools.client.atcoder import AtCoderClient, Contest, LoginError
+from atcodertools.client.atcoder import AtCoderClient, Contest, LoginError, PageNotFoundError
 from atcodertools.client.models.problem import Problem
 from atcodertools.client.models.problem_content import InputFormatDetectionError, SampleDetectionError
 from atcodertools.codegen.code_style_config import DEFAULT_WORKSPACE_DIR_PATH
@@ -169,17 +169,18 @@ def prepare_contest(atcoder_client: AtCoderClient,
                     retry_max_tries: int = 10):
     attempt_count = 1
     while True:
-        problem_list = atcoder_client.download_problem_list(
-            Contest(contest_id=contest_id))
-        if problem_list:
+        try:
+            problem_list = atcoder_client.download_problem_list(
+                Contest(contest_id=contest_id))
             break
-        if 0 < retry_max_tries < attempt_count:
-            raise EnvironmentInitializationError
-        logger.warning(
-            "Failed to fetch. Will retry in {} seconds. (Attempt {})".format(retry_delay_secs, attempt_count))
-        time.sleep(retry_delay_secs)
-        retry_delay_secs = min(retry_delay_secs * 2, retry_max_delay_secs)
-        attempt_count += 1
+        except PageNotFoundError:
+            if 0 < retry_max_tries < attempt_count:
+                raise EnvironmentInitializationError
+            logger.warning(
+                "Failed to fetch. Will retry in {} seconds. (Attempt {})".format(retry_delay_secs, attempt_count))
+            time.sleep(retry_delay_secs)
+            retry_delay_secs = min(retry_delay_secs * 2, retry_max_delay_secs)
+            attempt_count += 1
 
     tasks = [(atcoder_client,
               problem,