#318: add a function to get information of problems

onlinejudge/_implementation/utils.py

@@ -214,6 +214,11 @@ def is_update_available_on_pypi() -> bool:
     return a < b
 
 
+def remove_prefix(s: str, prefix: str) -> str:
+    assert s.startswith(prefix)
+    return s[len(prefix):]
+
+
 def remove_suffix(s: str, suffix: str) -> str:
     assert s.endswith(suffix)
     return s[:-len(suffix)]
@@ -229,9 +234,10 @@ def getter_with_load_details(name: str, check_with: Optional[str] = None) -> Cal
     @functools.wraps(lambda self: getattr(self, name))
     def wrapper(self, session: Optional[requests.Session] = None):
         if getattr(self, name) is None:
+            assert session is None or isinstance(session, requests.Session)
             getattr(self, '_load_details')(session)
         attr = getattr(self, name)
-        assert attr is not None
+        assert check_with is not None
         return attr
 
     return wrapper

onlinejudge/service/atcoder.py

@@ -157,6 +157,9 @@ class AtCoderContest(object):
     """
 
     def __init__(self, contest_id: str):
+        if contest_id.startswith('http'):
+            # an exception should be raised since mypy cannot check this kind of failure
+            raise ValueError('You should use AtCoderContest.from_url(url) instead of AtCoderContest(url)')
         self.contest_id = contest_id
 
         # NOTE: some fields remain undefined, comparing `_from_table_row`
@@ -166,6 +169,17 @@ def __init__(self, contest_id: str):
         self._duration_text = None  # type: Optional[str]
         self._rated_range = None  # type: Optional[str]
 
+    def get_url(self, type: Optional[str] = None, lang: Optional[str] = None) -> str:
+        if type is None or type == 'beta':
+            url = 'https://atcoder.jp/contests/{}'.format(self.contest_id)
+        elif type == 'old':
+            url = 'http://{}.contest.atcoder.jp/'.format(self.contest_id)
+        else:
+            assert False
+        if lang is not None:
+            url += '?lang={}'.format(lang)
+        return url
+
     @classmethod
     def from_url(cls, url: str) -> Optional['AtCoderContest']:
         """
@@ -278,6 +292,8 @@ def __init__(self, contest_id: str, problem_id: str):
         self._time_limit_msec = None  # type: Optional[int]
         self._memory_limit_byte = None  # type: Optional[int]
         self._alphabet = None  # type: Optional[str]
+        self._score = None  # type: Optional[int]
+        self._score_checked = None  # type: Optional[bool]
 
     @classmethod
     def _from_table_row(cls, tr: bs4.Tag) -> 'AtCoderProblem':
@@ -352,8 +368,16 @@ def _find_sample_tags(self, soup) -> Generator[Tuple[bs4.Tag, bs4.Tag], None, No
                     if prv and prv.name == 'h3' and prv.string:
                         yield (pre, prv)
 
-    def get_url(self) -> str:
-        return 'http://{}.contest.atcoder.jp/tasks/{}'.format(self.contest_id, self.problem_id)
+    def get_url(self, type: Optional[str] = None, lang: Optional[str] = None) -> str:
+        if type is None or type == 'beta':
+            url = 'https://atcoder.jp/contests/{}/tasks/{}'.format(self.contest_id, self.problem_id)
+        elif type == 'old':
+            url = 'http://{}.contest.atcoder.jp/tasks/{}'.format(self.contest_id, self.problem_id)
+        else:
+            assert False
+        if lang is not None:
+            url += '?lang={}'.format(lang)
+        return url
 
     def get_service(self) -> AtCoderService:
         return AtCoderService()
@@ -390,7 +414,7 @@ def from_url(cls, s: str) -> Optional['AtCoderProblem']:
     def get_input_format(self, session: Optional[requests.Session] = None) -> str:
         session = session or utils.new_default_session()
         # get
-        resp = _request('GET', self.get_url(), session=session)
+        resp = _request('GET', self.get_url(type='old'), session=session)
         msgs = AtCoderService._get_messages_from_cookie(resp.cookies)
         if AtCoderService._report_messages(msgs, unexpected=True):
             return ''
@@ -486,7 +510,7 @@ def _get_task_id(self, session: Optional[requests.Session] = None) -> int:
         if self._task_id is None:
             session = session or utils.new_default_session()
             # get
-            resp = _request('GET', self.get_url(), session=session)
+            resp = _request('GET', self.get_url(type='old'), session=session)
             msgs = AtCoderService._get_messages_from_cookie(resp.cookies)
             if AtCoderService._report_messages(msgs, unexpected=True):
                 raise SubmissionError
@@ -501,13 +525,30 @@ def _get_task_id(self, session: Optional[requests.Session] = None) -> int:
             self._task_id = int(m.group(1))
         return self._task_id
 
-    def _load_details(self, session: Optional[requests.Session] = None) -> int:
-        raise NotImplementedError
+    def _load_details(self, session: Optional[requests.Session] = None) -> None:
+        session = session or utils.new_default_session()
+
+        # get
+        resp = _request('GET', self.get_url(type='beta', lang='ja'), session=session)
+        soup = bs4.BeautifulSoup(resp.content.decode(resp.encoding), utils.html_parser)
+
+        # parse
+        h2 = soup.find('span', class_='h2')
+        self._alphabet, _, self._task_name = h2.text.partition(' - ')
+        time_limit, memory_limit = h2.find_next_sibling('p').text.split(' / ')
+        self._time_limit_msec = int(utils.remove_suffix(utils.remove_prefix(time_limit, '実行時間制限: '), ' sec')) * 1000
+        self._memory_limit_byte = int(utils.remove_suffix(utils.remove_prefix(memory_limit, 'メモリ制限: '), ' MB')) * 1000 * 1000
+        task_statement = soup.find('div', id='task-statement')
+        p = task_statement.find('p')  # first
+        if p is not None and p.text.startswith('配点 : '):
+            self._score = int(utils.remove_suffix(utils.remove_prefix(p.text, '配点 : '), ' 点'))
+        self._score_checked = True
 
     get_task_name = utils.getter_with_load_details('_task_name')  # type: Callable[..., str]
     get_time_limit_msec = utils.getter_with_load_details('_time_limit_msec')  # type: Callable[..., int]
     get_memory_limit_byte = utils.getter_with_load_details('_memory_limit_byte')  # type: Callable[..., int]
     get_alphabet = utils.getter_with_load_details('_alphabet')  # type: Callable[..., str]
+    get_score = utils.getter_with_load_details('_score', check_with='_score_checked')  # type: Callable[..., str]
 
 
 class AtCoderSubmission(onlinejudge.type.Submission):
@@ -569,8 +610,8 @@ def from_url(cls, s: str, problem_id: Optional[str] = None) -> Optional['AtCoder
 
         return None
 
-    def get_url(self, type: str = 'latest', lang: Optional[str] = None) -> str:
-        if type == 'latest' or type == 'beta':
+    def get_url(self, type: Optional[str] = None, lang: Optional[str] = None) -> str:
+        if type is None or type == 'beta':
             url = 'https://atcoder.jp/contests/{}/submissions/{}'.format(self.contest_id, self.submission_id)
         elif type == 'old':
             url = 'https://{}.contest.atcoder.jp/submissions/{}'.format(self.contest_id, self.submission_id)

tests/service_atcoder.py

@@ -35,7 +35,7 @@ def test_from_url(self):
         self.assertEqual(AtCoderContest.from_url('https://atcoder.jp/contests/'), None)
 
     def test_list_problems(self):
-        contest = AtCoderContest('agc028')
+        contest = AtCoderContest.from_url('https://atcoder.jp/contests/agc028')
         problems = contest.list_problems()
         self.assertEqual(len(problems), 7)
         self.assertEqual(problems[0].get_alphabet(), 'A')
@@ -55,6 +55,22 @@ def test_from_url(self):
         self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/agc030/tasks/agc030_c').contest_id, 'agc030')
         self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/agc030/tasks/agc030_c').problem_id, 'agc030_c')
 
+    def test_load_details(self):
+        problem = AtCoderProblem.from_url('https://atcoder.jp/contests/abc118/tasks/abc118_a')
+        self.assertEqual(problem.get_alphabet(), 'A')
+        self.assertEqual(problem.get_task_name(), 'B +/- A')
+        self.assertEqual(problem.get_time_limit_msec(), 2000)
+        self.assertEqual(problem.get_memory_limit_byte(), 1024 * 1000 * 1000)
+        self.assertEqual(problem.get_score(), 100)
+
+    def test_get_alphabet(self):
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/agc028/tasks/agc028_f').get_alphabet(), 'F')
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/agc028/tasks/agc028_f2').get_alphabet(), 'F2')
+
+    def test_get_score(self):
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/future-contest-2018-final/tasks/future_contest_2018_final_a').get_score(), 50000000)
+        self.assertEqual(AtCoderProblem.from_url('https://atcoder.jp/contests/abc001/tasks/abc001_4').get_score(), None)
+
 
 class AtCoderSubmissionTest(unittest.TestCase):
     def test_from_url(self):