find_original_update_blocks now accepts fence param, raises on mangle…

…d <source>filename.ext #317

aider/coders/editblock_coder.py

@@ -17,15 +17,15 @@ def get_edits(self):
         content = self.partial_response_content
 
         # might raise ValueError for malformed ORIG/UPD blocks
-        edits = list(find_original_update_blocks(content))
+        edits = list(find_original_update_blocks(content, self.fence))
 
         return edits
 
     def apply_edits(self, edits):
         for path, original, updated in edits:
             full_path = self.abs_root_path(path)
             content = self.io.read_text(full_path)
-            content = do_replace(full_path, content, original, updated)
+            content = do_replace(full_path, content, original, updated, self.fence)
             if content:
                 self.io.write_text(full_path, content)
                 continue
@@ -247,7 +247,10 @@ def replace_closest_edit_distance(whole_lines, part, part_lines, replace_lines):
     return modified_whole
 
 
-def strip_quoted_wrapping(res, fname=None, fence=None):
+DEFAULT_FENCE = ("`" * 3, "`" * 3)
+
+
+def strip_quoted_wrapping(res, fname=None, fence=DEFAULT_FENCE):
     """
     Given an input string which may have extra "wrapping" around it, remove the wrapping.
     For example:
@@ -261,9 +264,6 @@ def strip_quoted_wrapping(res, fname=None, fence=None):
     if not res:
         return res
 
-    if not fence:
-        fence = ("```", "```")
-
     res = res.splitlines()
 
     if fname and res[0].strip().endswith(Path(fname).name):
@@ -310,7 +310,23 @@ def do_replace(fname, content, before_text, after_text, fence=None):
 split_re = re.compile(r"^((?:" + separators + r")[ ]*\n)", re.MULTILINE | re.DOTALL)
 
 
-def find_original_update_blocks(content):
+missing_filename_err = f"Bad/missing filename. Filename should be alone on the line before {HEAD}"
+
+
+def strip_filename(filename, fence):
+    filename = filename.strip()
+
+    if filename == "...":
+        return
+
+    start_fence = fence[0]
+    if filename.startswith(start_fence):
+        return
+
+    return filename
+
+
+def find_original_update_blocks(content, fence=DEFAULT_FENCE):
     # make sure we end with a newline, otherwise the regex will miss <<UPD on the last line
     if not content.endswith("\n"):
         content = content + "\n"
@@ -337,22 +353,20 @@ def find_original_update_blocks(content):
 
             processed.append(cur)  # original_marker
 
-            filename = processed[-2].splitlines()[-1].strip()
+            filename = strip_filename(processed[-2].splitlines()[-1], fence)
             try:
-                if not len(filename) or "`" in filename:
-                    filename = processed[-2].splitlines()[-2].strip()
-                if not len(filename) or "`" in filename or filename == "...":
+                if not filename:
+                    filename = strip_filename(processed[-2].splitlines()[-2], fence)
+                if not filename:
                     if current_filename:
                         filename = current_filename
                     else:
-                        raise ValueError(
-                            f"Bad/missing filename. It should go right above the {HEAD}"
-                        )
+                        raise ValueError(missing_filename_err)
             except IndexError:
                 if current_filename:
                     filename = current_filename
                 else:
-                    raise ValueError(f"Bad/missing filename. It should go right above the {HEAD}")
+                    raise ValueError(missing_filename_err)
 
             current_filename = filename
 

tests/test_editblock.py

@@ -80,6 +80,29 @@ def test_find_original_update_blocks(self):
         edits = list(eb.find_original_update_blocks(edit))
         self.assertEqual(edits, [("foo.txt", "Two\n", "Tooooo\n")])
 
+    def test_find_original_update_blocks_mangled_filename_w_source_tag(self):
+        source = "source"
+
+        edit = """
+Here's the change:
+
+<%s>foo.txt
+<<<<<<< SEARCH
+One
+=======
+Two
+>>>>>>> REPLACE
+</%s>
+
+Hope you like it!
+""" % (source, source)
+
+        fence = ("<%s>" % source, "</%s>" % source)
+
+        with self.assertRaises(ValueError) as cm:
+            _edits = list(eb.find_original_update_blocks(edit, fence))
+        self.assertIn("missing filename", str(cm.exception))
+
     def test_find_original_update_blocks_quote_below_filename(self):
         edit = """
 Here's the change: