code optimization

Classifier/collect.py

@@ -24,17 +24,15 @@ def choose_random_wrapper(config_code, year):
 def choose_random(db_repo, config_code, year, save_dir, max_commits):
     repo_node = get_config_nodes_repo_dict()[config_code]
     commit_utils = CommitUtils(False, repo_node, None)
-    vccs = db_repo.get_all_vccs()
+    vccs = set(db_repo.get_all_vccs())
     commits = commit_utils.get_commits_between_years(year, year)
     commits = list(filter(lambda com: com.hexsha not in vccs, commits))
-    already_saved = glob.glob(f"{save_dir}/*.json")
-    i = 0
 
-    for saved in already_saved:
-        saved = saved.split("/")[-1].rstrip(".json")
-        commit_shas = [commit.hexsha for commit in commits]
-        if saved in commit_shas:
-            i += 1
+    already_saved = set(glob.glob(f"{save_dir}/*.json"))
+    already_saved = {os.path.basename(f).rstrip(".json") for f in already_saved}
+    commit_shas = {commit.hexsha for commit in commits}
+
+    i = len(already_saved.intersection(commit_shas))
 
     for _ in commits:
         try:
@@ -42,16 +40,19 @@ def choose_random(db_repo, config_code, year, save_dir, max_commits):
                 print(f"PID {os.getpid()}: Saved {i} commits for {config_code} in {year}")
                 break
             random_commit = random.choice(commits)
-            if os.path.exists(f"{save_dir}/{random_commit.hexsha}.json"):
+            if random_commit.hexsha in already_saved:
                 continue
             commit_features = CommitFeatures(repo_node.find("./path").text, random_commit.hexsha)
             commit_features.extract_features()
             commit_features.save_features_to_json(save_dir)
             db_repo.save_commit(random_commit, config_code)
             i += 1
+            already_saved.add(random_commit.hexsha)
             print(f"PID {os.getpid()}: Saved {random_commit.hexsha} for {config_code} in {year}")
         except ValueError as e:
             print(e)
+        except Exception as e:
+            print(f"PID {os.getpid()}: Error while saving commit for {config_code} in {year}")
 
 
 def store_unclassified_training():