Add problems_v2 method

RussellDash332 · Feb 12, 2024 · f801869 · f801869
1 parent ad3590f
commit f801869
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 29 deletions.
diff --git a/README.md b/README.md
@@ -39,6 +39,9 @@ kt.problems()                               # problems you have solved so far
 kt.problems(show_partial=False)             # exclude partial submissions
 kt.problems(*[True]*4)                      # literally all problems on Kattis
 
+kt.problems_v2()                            # problems you have solved so far
+kt.problems_v2(show_non_ac=False)           # literally all problems on Kattis
+
 kt.list_unsolved()                          # let's grind!
 
 kt.plot_problems()                          # plot the points distribution

diff --git a/autokattis/api/__init__.py b/autokattis/api/__init__.py
@@ -170,6 +170,74 @@ def problems(self, show_solved=True, show_partial=True, show_tried=False, show_u
                             })
         return self.Result(sorted(data, key=lambda x: x['id']))
 
+    @lru_cache
+    def problems_v2(self, show_non_ac=False):
+        '''
+        Gets all accepted Kattis problems. Note that this is entirely different than the `problems`
+            method due to possible gateway issues from the initial version.
+        Returns a simpler JSON-like structure with these fields:
+            name, id, link
+
+        Default: all solved and partially solved problems.
+        '''
+
+        has_content = True
+        params = {
+            'page': 0,
+            'tab': 'submissions',
+            'status': 'AC'
+        }
+
+        data = []
+        pid_set = set()
+
+        with ThreadPoolExecutor(max_workers=self.MAX_WORKERS) as executor:
+            futures = []
+            while has_content:
+                has_content = False
+                futures.clear()
+                for _ in range(self.MAX_WORKERS):
+                    futures.append(executor.submit(self.new_get, f'{self.BASE_URL}/users/{self.user}', params=params.copy()))
+                    params['page'] += 1
+                for f in as_completed(futures):
+                    response = f.result()
+                    soup = bs(response.content, features='lxml')
+                    if not soup: continue
+                    table = table = soup.find('div', id='submissions-tab').find('section', class_='strip strip-item-plain').find('table', class_='table2')
+                    try:
+                        table_content = table.tbody.find_all('tr')
+                    except AttributeError:
+                        continue
+                    for row in table_content:
+                        columns = row.find_all('td')
+                        if columns and len(columns) > 2:
+                            has_content = True
+                            pid = columns[2].find_all('a')[-1].get('href').split('/')[-1] # might have two links if it belongs to a contest
+                            if pid not in pid_set:
+                                link = f"{self.BASE_URL}/problems/{pid}"
+                                name = columns[2].text.split('/')[-1].strip()
+                                pid_set.add(pid)
+                                data.append({
+                                    'name': name,
+                                    'id': pid,
+                                    'link': link
+                                })
+
+        if show_non_ac:
+            # we can just use the latest soup and take the dropdown
+            for option in soup.find_all('option')[1:]:
+                pid = option.get('value').strip()
+                if not pid: break
+                if pid not in pid_set:
+                    pid_set.add(pid)
+                    data.append({
+                        'name': option.text.strip(),
+                        'id': pid,
+                        'link': f"{self.BASE_URL}/problems/{pid}"
+                    })
+
+        return self.Result(sorted(data, key=lambda x: x['id']))
+
     @lru_cache
     def plot_problems(self, filepath=None, show_solved=True, show_partial=True, show_tried=False, show_untried=False):
         '''

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup_args = dict(
     name='autokattis',
-    version='1.6.3',
+    version='1.6.4',
     description='Updated Kattis API wrapper',
     long_description_content_type="text/markdown",
     long_description=README,

diff --git a/test/main.py b/test/main.py
@@ -8,6 +8,11 @@
 print(df:=ret.to_df())
 df.to_csv('test_problems_default.csv', index=False)
 
+print('=== TEST PROBLEMS V2 (DEFAULT) ===')
+ret = kt.problems_v2()                              # problems you have solved so far
+print(df:=ret.to_df())
+df.to_csv('test_problems_v2_default.csv', index=False)
+
 print('=== TEST PROBLEMS (NO PARTIAL) ===')
 ret = kt.problems(show_partial=False)               # exclude partial submissions
 print(df:=ret.to_df())
@@ -18,6 +23,11 @@
 print(df:=ret.to_df())
 df.to_csv('test_problems_all.csv', index=False)
 
+print('=== TEST PROBLEMS V2 (ALL PROBLEMS) ===')
+ret = kt.problems_v2(show_non_ac=True)              # literally all problems on Kattis
+print(df:=ret.to_df())
+df.to_csv('test_problems_v2_all.csv', index=False)
+
 print('=== TEST LIST UNSOLVED ===')
 ret = kt.list_unsolved()                            # let's grind!
 print(df:=ret.to_df())

diff --git a/test/nus.py b/test/nus.py
@@ -3,37 +3,15 @@
 
 kt = NUSKattis(USER, PASSWORD)
 
-print('=== TEST PROBLEMS (DEFAULT) ===')
-ret = kt.problems()                                 # problems you have solved so far
+print('=== TEST PROBLEMS V2 (DEFAULT) ===')
+ret = kt.problems_v2()                              # problems you have solved so far
 print(df:=ret.to_df())
-df.to_csv('test_nus_problems_default.csv', index=False)
+df.to_csv('test_nus_problems_v2_default.csv', index=False)
 
-print('=== TEST PROBLEMS (NO PARTIAL) ===')
-ret = kt.problems(show_partial=False)               # exclude partial submissions
+print('=== TEST PROBLEMS V2 (ALL PROBLEMS) ===')
+ret = kt.problems_v2(show_non_ac=True)              # literally all problems on Kattis
 print(df:=ret.to_df())
-df.to_csv('test_nus_problems_no_partial.csv', index=False)
-
-print('=== TEST PROBLEMS (ALL PROBLEMS) ===')
-ret = kt.problems(*[True]*4)                        # literally all problems on Kattis
-print(df:=ret.to_df())
-df.to_csv('test_nus_problems_all.csv', index=False)
-
-print('=== TEST LIST UNSOLVED ===')
-ret = kt.list_unsolved()                            # let's grind!
-print(df:=ret.to_df())
-df.to_csv('test_nus_list_unsolved.csv', index=False)
-
-print('=== TEST PLOT PROBLEMS (DEFAULT) ===')
-ret = kt.plot_problems()                            # plot the points distribution
-print('Done!')
-
-print('=== TEST PLOT PROBLEMS (TO FILEPATH) ===')
-ret = kt.plot_problems(filepath='nus_plot.png')     # save to a filepath
-print('Done!')
-
-print('=== TEST PLOT PROBLEMS (NO PARTIAL) ===')
-ret = kt.plot_problems(show_partial=False)          # plot fully solved submissions
-print('Done!')
+df.to_csv('test_nus_problems_v2_all.csv', index=False)
 
 print('=== TEST PROBLEM (SINGLE) ===')
 ret = kt.problem('2048')                            # fetch info about a problem