Seperate out find_descriptors

Fast-64 · Sep 27, 2024 · c2eebb9 · c2eebb9
1 parent 82f7e22
commit c2eebb9
Showing 1 changed file with 89 additions and 73 deletions.
diff --git a/fast64_internal/sm64/sm64_utility.py b/fast64_internal/sm64/sm64_utility.py
@@ -1,4 +1,4 @@
-from typing import NamedTuple
+from typing import NamedTuple, Optional
 from pathlib import Path
 import os
 import re
@@ -140,43 +140,97 @@ def __init__(self, string: str, regex: str = ""):
             self.regex = re.escape(string) + r"\n?"
 
 
+class DescriptorMatch(NamedTuple):
+    string: str
+    start: int
+    end: int
+
+
+def find_descriptor_in_text(value: ModifyFoundDescriptor, commentless: str, comment_map: dict):
+    matches: list[DescriptorMatch] = []
+    for match in re.finditer(value.regex, commentless):
+        found_start, found_end = match.start(), match.end()
+        start, end = match.start(), match.end()
+        for commentless_start, comment_size in comment_map:  # only remove parts outside comments
+            if commentless_start >= found_start and commentless_start <= found_end:  # comment inside descriptor
+                end += comment_size
+            elif found_end >= commentless_start:
+                start += comment_size
+                end += comment_size
+        matches.append(DescriptorMatch(match.group(0), start, end))
+    return matches
+
+
+def find_descriptors(
+    text: str,
+    descriptors: list[ModifyFoundDescriptor],
+    error_if_no_header=False,
+    header: Optional[ModifyFoundDescriptor] = None,
+    error_if_no_footer=False,
+    footer: Optional[ModifyFoundDescriptor] = None,
+    ignore_comments=True,
+):
+    """Returns: The found matches from descriptors, the footer pos (the end of the text if none)"""
+    comment_map = []
+    if ignore_comments:
+        commentless = ""
+        last_pos = 0
+
+        for match in re.finditer(COMMENT_PATTERN, text):
+            commentless += text[last_pos : match.start()]  # add text before comment
+            string = match.group(0)
+            if string.startswith("/"):
+                comment_map.append((len(commentless), len(string) - 1))
+                commentless += " "
+            else:
+                commentless += string
+            last_pos = match.end()
+
+        commentless += text[last_pos:]  # add any remaining text after the last match
+    else:
+        commentless = text
+
+    header_matches = find_descriptor_in_text(header, commentless, comment_map) if header is not None else []
+    footer_matches = find_descriptor_in_text(footer, commentless, comment_map) if footer is not None else []
+
+    header_pos = 0
+    if len(header_matches) > 0:
+        _, header_pos, _ = header_matches[0]
+    elif error_if_no_header:
+        raise PluginError(f"Header {header.string} does not exist.")
+
+    # find first footer after the header
+    if footer_matches:
+        if header_matches:
+            footer_pos = next((pos for _, pos, _ in footer_matches if pos >= header_pos), footer_matches[-1])
+        elif len(footer_matches) > 0:
+            _, footer_pos, _ = footer_matches[-1]
+    else:
+        if error_if_no_footer:
+            raise PluginError(f"Footer {footer.string} does not exist.")
+        footer_pos = len(text)
+
+    found_matches: dict[ModifyFoundDescriptor, list[DescriptorMatch]] = {}
+    for descriptor in descriptors:
+        matches = find_descriptor_in_text(descriptor, commentless[header_pos:footer_pos], comment_map)
+        if matches:
+            found_matches.setdefault(descriptor, []).extend(matches)
+    return found_matches, footer_pos
+
+
 def write_or_delete_if_found(
     path: Path,
-    to_add: list[ModifyFoundDescriptor] | None = None,
-    to_remove: list[ModifyFoundDescriptor] | None = None,
+    to_add: Optional[list[ModifyFoundDescriptor]] = None,
+    to_remove: Optional[list[ModifyFoundDescriptor]] = None,
     path_must_exist=False,
     create_new=False,
     error_if_no_header=False,
-    header: ModifyFoundDescriptor | None = None,
+    header: Optional[ModifyFoundDescriptor] = None,
     error_if_no_footer=False,
-    footer: ModifyFoundDescriptor | None = None,
+    footer: Optional[ModifyFoundDescriptor] = None,
     ignore_comments=True,
 ):
-    class DescriptorMatch(NamedTuple):
-        string: str
-        start: int
-        end: int
-
-    def find_descriptor_in_text(value: ModifyFoundDescriptor, commentless: str, comment_map: dict):
-        matches: list[DescriptorMatch] = list()
-        for match in re.finditer(value.regex, commentless):
-            found_start, found_end = match.start(), match.end()
-            start, end = match.start(), match.end()
-            for commentless_start, comment_size in comment_map:  # only remove parts outside comments
-                if commentless_start >= found_start and commentless_start <= found_end:  # comment inside descriptor
-                    end += comment_size
-                elif found_end >= commentless_start:
-                    start += comment_size
-                    end += comment_size
-            matches.append(DescriptorMatch(match.group(0), start, end))
-        return matches
-
-    comment_map = []
-    commentless = ""
-    found_matches: dict[ModifyFoundDescriptor, list[DescriptorMatch]] = {}
     changed = False
-    text = ""
-    header_pos = footer_pos = 0
     to_add, to_remove = to_add or [], to_remove or []
 
     assert not (path_must_exist and create_new), "path_must_exist and create_new"
@@ -187,51 +241,13 @@ def find_descriptor_in_text(value: ModifyFoundDescriptor, commentless: str, comm
 
     if os.path.exists(path) and not create_new:
         text = path.read_text()
-
-        if ignore_comments:
-            commentless = ""
-            last_pos = 0
-
-            for match in re.finditer(COMMENT_PATTERN, text):
-                commentless += text[last_pos : match.start()]  # add text before comment
-                string = match.group(0)
-                if string.startswith("/"):
-                    comment_map.append((len(commentless), len(string) - 1))
-                    commentless += " "
-                else:
-                    commentless += string
-                last_pos = match.end()
-
-            commentless += text[last_pos:]  # add any remaining text after the last match
-        else:
-            commentless = text
-        if commentless and commentless[-1] not in {"\n", "\r"}:  # add end new line if not there
+        if text and text[-1] not in {"\n", "\r"}:  # add end new line if not there
             text += "\n"
-
-        header_matches = find_descriptor_in_text(header, commentless, comment_map) if header is not None else []
-        footer_matches = find_descriptor_in_text(footer, commentless, comment_map) if footer is not None else []
-
-        header_pos = 0
-        if len(header_matches) > 0:
-            _, header_pos, _ = header_matches[0]
-        elif error_if_no_header:
-            raise PluginError(f"Header {header.string} does not exist.")
-
-        # find first footer after the header
-        if footer_matches:
-            if header_matches:
-                footer_pos = next((pos for _, pos, _ in footer_matches if pos >= header_pos), footer_matches[-1])
-            elif len(footer_matches) > 0:
-                _, footer_pos, _ = footer_matches[-1]
-        else:
-            if error_if_no_footer:
-                raise PluginError(f"Footer {footer.string} does not exist.")
-            footer_pos = len(text)
-
-        for descriptor in to_add + to_remove:
-            matches = find_descriptor_in_text(descriptor, commentless[header_pos:footer_pos], comment_map)
-            if matches:
-                found_matches.setdefault(descriptor, []).extend(matches)
+        found_matches, footer_pos = find_descriptors(
+            text, to_add + to_remove, error_if_no_header, header, error_if_no_footer, footer, ignore_comments
+        )
+    else:
+        text, found_matches, footer_pos = "", {}, 0
 
     for descriptor in to_remove:
         matches = found_matches.get(descriptor)