Refactoring, cleaning

Author: trsvchn

src/entrypoint.py

@@ -7,6 +7,9 @@
 # Badge setup
 SRC = '"https://colab.research.google.com/assets/colab-badge.svg"'
 ALT = '"Open In Colab"'
+BADGE_VAR = "{{ badge }}"
+BADGE_PATTERN = re.compile(r"<!--<badge>-->(.*?)<!--</badge>-->")
+HREF_PATTERN = re.compile(r"href=[\"\'](.*?)[\"\']")
 
 # Set repository
 CURRENT_REPOSITORY = os.environ['GITHUB_REPOSITORY']
@@ -77,12 +80,20 @@ def check_cells(data: dict, repo_name: str, branch: str, nb_path: str):
     for cell in data["cells"]:
         # Check only markdown cells
         if cell["cell_type"] == "markdown":
-            # If a cell already has a badge - check the repo and branch
-            if UPDATE and any(["<!--<badge>-->" in line for line in cell["source"]]):
-                # Update repo, branch, file path
-                save = True if update_badge(cell, repo_name, branch, nb_path) else save
-            # Add badge code, add metadata
-            save = True if add_badge(cell, repo_name, branch, nb_path) else save
+            text = cell["source"]
+            for i, line in enumerate(text):
+                # If a cell already has a badge - check the repo and branch
+                if UPDATE:
+                    # Update repo, branch, file path
+                    new_line = update_badge(line, repo_name, branch, nb_path)
+                    if new_line:
+                        text[i] = new_line
+                        save = True if new_line else save
+                # Add badge code
+                new_line = add_badge(line, repo_name, branch, nb_path)
+                if new_line:
+                    text[i] = new_line
+                    save = True if new_line else save
         else:
             continue
 
@@ -92,22 +103,15 @@ def check_cells(data: dict, repo_name: str, branch: str, nb_path: str):
         write_nb(data, nb_path)
 
 
-def add_badge(cell: dict, repo_name: str, branch: str, nb_path: str):
+def add_badge(line: str, repo_name: str, branch: str, nb_path: str):
     """Inserts "Open in Colab" badge.
     """
-    modified = False
-    pattern = "{{ badge }}"
-    text = cell["source"]
-    for i, line in enumerate(text):
-        if pattern in line:
-            badge = prepare_badge_code(repo_name, branch, nb_path)
-            print(f"{nb_path}: Inserting badge...")
-            new_line = line.replace(pattern, badge)
-            text[i] = new_line
-            modified = True
-        else:
-            continue
-    return modified
+    new_line = None
+    if BADGE_VAR in line:
+        badge = prepare_badge_code(repo_name, branch, nb_path)
+        print(f"{nb_path}: Inserting badge...")
+        new_line = line.replace(BADGE_VAR, badge)
+    return new_line
 
 
 def prepare_badge_code(repo_name: str, branch: str, nb_path: str) -> str:
@@ -118,33 +122,25 @@ def prepare_badge_code(repo_name: str, branch: str, nb_path: str) -> str:
     return code
 
 
-def update_badge(cell: dict, repo_name: str, branch: str, nb_path: str):
+def update_badge(line: str, repo_name: str, branch: str, nb_path: str):
     """Updates added badge code.
     """
-    modified = False
-    badge_pattern = re.compile(r"<!--<badge>-->(.*?)<!--</badge>-->")
-    href_pattern = re.compile(r"href=[\"\'](.*?)[\"\']")
-
     new_href = f"https://colab.research.google.com/github/{repo_name}/blob/{branch}/{nb_path}"
-    text = cell["source"]
-
-    for i, line in enumerate(text):
-        badges = badge_pattern.findall(line)
-        if badges:
-            for badge in badges:
-                href = href_pattern.findall(badge)[0]
-                repo_branch_nb = href.split("/github/")[-1]
-                curr_repo, branch_nb = repo_branch_nb.split("/blob/")
-                curr_branch, curr_nb_path = branch_nb.split("/", 1)
-
-                if (curr_repo != repo_name) or (curr_branch != branch) or (curr_nb_path != nb_path):
-                    print(f"{nb_path}: Updating badge info...")
-                    new_line = line.replace(href, new_href)
-                    text[i] = new_line
-                    modified = True
-        else:
-            continue
-    return modified
+    new_line = None
+
+    badges = BADGE_PATTERN.findall(line)
+    if badges:
+        for badge in badges:
+            href = HREF_PATTERN.findall(badge)[0]
+            repo_branch_nb = href.split("/github/")[-1]
+            curr_repo, branch_nb = repo_branch_nb.split("/blob/")
+            curr_branch, curr_nb_path = branch_nb.split("/", 1)
+
+            if (curr_repo != repo_name) or (curr_branch != branch) or (curr_nb_path != nb_path):
+                print(f"{nb_path}: Updating badge info...")
+                new_line = line.replace(href, new_href)
+
+    return new_line
 
 
 def write_nb(data: dict, file_path: str) -> None: