diff --git a/fast64_internal/sm64/sm64_utility.py b/fast64_internal/sm64/sm64_utility.py index 03be9d5a9..bbd590e5a 100644 --- a/fast64_internal/sm64/sm64_utility.py +++ b/fast64_internal/sm64/sm64_utility.py @@ -146,9 +146,11 @@ def write_or_delete_if_found( to_remove: list[ModifyFoundDescriptor] | None = None, path_must_exist=False, create_new=False, + error_if_no_header=False, + header: ModifyFoundDescriptor | None = None, + error_if_no_footer=False, footer: ModifyFoundDescriptor | None = None, ignore_comments=True, - error_if_no_footer=False, ): class DescriptorMatch(NamedTuple): string: str @@ -173,7 +175,7 @@ def find_descriptor_in_text(value: ModifyFoundDescriptor, commentless: str, comm found_matches: dict[ModifyFoundDescriptor, list[DescriptorMatch]] = {} changed = False text = "" - footer_pos = 0 + 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" @@ -205,16 +207,24 @@ def find_descriptor_in_text(value: ModifyFoundDescriptor, commentless: str, comm if commentless and commentless[-1] not in {"\n", "\r"}: # add end new line if not there text += "\n" - footer_pos = len(text) - if footer is not None: # Find footer, optionally error if not found - footer_matches = find_descriptor_in_text(footer, commentless, comment_map) - if len(footer_matches) > 0: - _, footer_pos, _ = footer_matches[-1] - elif error_if_no_footer: + 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[-1] + elif error_if_no_header: + raise PluginError(f"Header {header.string} does not exist.") + + # find first footer after the header + footer_pos = next((pos for _, pos, _ in footer_matches if pos >= header_pos), None) + if footer_pos is None: + 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[:footer_pos], comment_map) + matches = find_descriptor_in_text(descriptor, commentless[header_pos:footer_pos], comment_map) if matches: found_matches.setdefault(descriptor, []).extend(matches)