Update arvix.py

arvix.py

@@ -1,15 +1,19 @@
 import json
 import random
-import tools
+import os
+import sys
 from bs4 import BeautifulSoup
 
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
+import utils
+
 def fetch_new_page(category):
     url = f'https://arxiv.org/list/{category}/new'
-    return tools.fetch_page(url)
+    return utils.fetch_page(url)
 
 def fetch_recent_page(category):
     url = f'https://arxiv.org/list/{category}/recent'
-    return tools.fetch_page(url)
+    return utils.fetch_page(url)
 
 def extract_new_data(category):
     paper_ids = []
@@ -20,15 +24,12 @@ def extract_new_data(category):
         paper_contents = list.find_all('dd')
         titles = [paper_content.find('div', class_='list-title').text.strip().split('Title:')[-1].strip() for paper_content in paper_contents]
         for paper, title in zip(papers, titles):
-            if not tools.verify_simple_title(title):
+            if not utils.verify_simple_title(title):
                 continue
-            else:
-                paper_link = paper.find('a', href=True)
-                if paper_link:
-                    paper_id = paper_link.text.strip().split(':')[1]
-                    paper_ids.append(paper_id)
-                else:
-                    continue
+            paper_link = paper.find('a', href=True)
+            if paper_link:
+                paper_id = paper_link.text.strip().split(':')[1]
+                paper_ids.append(paper_id)
     return paper_ids
 
 def extract_recent_data(category):
@@ -42,42 +43,22 @@ def extract_recent_data(category):
             if paper_link:
                 paper_id = paper_link.text.strip().split(':')[1]
                 paper_ids.append(paper_id)
-            else:
-                continue
     return paper_ids
 
 def extract_data(category):
-    sanitized_data = []
+    all_ids = set()
     new_data = extract_new_data(category)
     recent_data = extract_recent_data(category)
-    data = list(set(new_data + recent_data))
-    if category in ["hep-ex", "hep-lat", "hep-ph", "hep-th"]:
-        category_list = []
-        for id in data:
-            if len(category_list) >= 1:
-                break
-            if tools.check_data_in_file(id, 'arxiv.txt'):
-                continue
-            else:
-                category_list.append(id)
-        for category_id in category_list:
-            sanitized_data.append(category_id)
-            tools.write_data_to_file(id, 'arxiv.txt')
-    else:
-        for id in data:
-            if len(sanitized_data) >= 3:
-                break
-            if tools.check_data_in_file(id, 'arxiv.txt'):
-                continue
-            else:
-                tools.write_data_to_file(id, 'arxiv.txt')
-                sanitized_data.append(id)
-    random.shuffle(sanitized_data)
-    return sanitized_data
+    combined_data = new_data + recent_data
+    for paper_id in combined_data:
+        if not utils.check_data_in_file(paper_id, 'arxiv.txt'):
+            utils.write_data_to_file(paper_id, 'arxiv.txt')
+            all_ids.add(paper_id)
+        if len(all_ids) >= 4:
+            break
+    return list(all_ids)
 
 def extract_arxiv_data():
-    if not tools.download_datafile('arxiv.txt'):
-        raise Exception("Failed to download datafile")
     categories = {
         "Astrophysics": ["astro-ph"],
         "Condensed Matter": ["cond-mat"],
@@ -98,33 +79,30 @@ def extract_arxiv_data():
         "Economics": ["econ"]
     }
     data = {}
+    used_ids = set()
+
     for category, subcategories in categories.items():
-        category_data = {}
-        all_ids = []
-        temp_id_storage = []
+        category_ids = set()
         for subcategory in subcategories:
             ids = extract_data(subcategory)
-            if len(ids) == 3:
-                for id in ids:
-                    temp_id_storage.append(id)
-            else:
-                for id in ids:
-                    all_ids.append(id)
-        for temp_id in temp_id_storage:
-            all_ids.append(temp_id)
-        random.shuffle(all_ids)
-        if len(all_ids) > 3:
-            print(f"Found more than 3 papers for {category}.")
-            all_ids = all_ids[:3]
-        category_data['count'] = len(all_ids)
-        category_data['ids'] = all_ids
-        data[category] = category_data
-    data = json.dumps(data, indent=4, ensure_ascii=False)
-    if not tools.upload_datafile('arxiv.txt'):
-        raise Exception("Failed to upload datafile")
-    return data
+            for paper_id in ids:
+                if paper_id not in used_ids:
+                    category_ids.add(paper_id)
+                    used_ids.add(paper_id)
+                if len(category_ids) == 4:
+                    break
+            if len(category_ids) == 4:
+                break
+
+        # Ensure exactly 4 IDs for each category
+        while len(category_ids) < 4:
+            category_ids.add(random.choice(list(used_ids)))
+
+        data[category] = {"ids": list(category_ids)}
+
+    return json.dumps(data, indent=4, ensure_ascii=False)
 
 if __name__ == '__main__':
     data = extract_arxiv_data()
     with open('arxiv_data.json', 'w') as f:
-        f.write(data)
+        f.write(data)