fix: fixed issues with applying patches

Author: ErikBjare

gptme/tools/patch.py

@@ -55,9 +55,9 @@
 ```
 """.strip()
 
-ORIGINAL = "\n<<<<<<< ORIGINAL\n"
+ORIGINAL = "<<<<<<< ORIGINAL\n"
 DIVIDER = "\n=======\n"
-UPDATED = ">>>>>>> UPDATED\n"
+UPDATED = "\n>>>>>>> UPDATED"
 
 
 def apply(codeblock: str, content: str) -> str:
@@ -77,9 +77,9 @@ def apply(codeblock: str, content: str) -> str:
         raise ValueError(f"invalid patch, no `{DIVIDER.strip()}`", codeblock)
     original, modified = re.split(DIVIDER, original)
 
-    if UPDATED not in modified:  # pragma: no cover
+    if UPDATED not in "\n" + modified:  # pragma: no cover
         raise ValueError(f"invalid patch, no `{UPDATED.strip()}`", codeblock)
-    modified = re.split(UPDATED, modified)[0].rstrip("\n")
+    modified = re.split(UPDATED, modified)[0]
 
     # TODO: maybe allow modified chunk to contain "// ..." to refer to chunks in the original,
     #       and then replace these with the original chunks?
@@ -93,12 +93,6 @@ def apply(codeblock: str, content: str) -> str:
 
 
 def apply_file(codeblock, filename):
-    codeblock = codeblock.strip()
-    _patch, filename = codeblock.splitlines()[0].split()
-    if not _patch == "```patch":
-        raise ValueError(
-            "invalid patch, codeblock is missing leading ```patch", codeblock
-        )
     if not Path(filename).exists():
         raise FileNotFoundError(filename)
 

tests/test_tools_patch.py

@@ -47,13 +47,25 @@ def hello(name="world"):
 """
     )
 
+
+def test_clear_file():
     # only remove code in patch
+    content = """
+def hello():
+    print("hello")
+
+if __name__ == "__main__":
+    hello()
+"""
+
+    # NOTE: test fails if UPDATED block doesn't have an empty line
     codeblock = """
 ```patch test.py
 <<<<<<< ORIGINAL
 def hello():
     print("hello")
 =======
+
 >>>>>>> UPDATED
 ```
 """
@@ -64,3 +76,27 @@ def hello():
     assert result.startswith(
         "\n\n"
     ), f"result: {result.replace(newline, newline_escape)}"
+
+
+def test_apply_empty_lines():
+    # a test where it replaces a empty line with 3 empty lines
+    # checks that whitespace is preserved
+    content = """
+def hello():
+    print("hello")
+
+if __name__ == "__main__":
+    hello()
+"""
+    codeblock = """
+```patch test.py
+<<<<<<< ORIGINAL
+
+=======
+
+
+>>>>>>> UPDATED
+```
+"""
+    result = apply(codeblock, content)
+    assert "\n\n\n" in result