diff --git a/aider/coders/wholefile_coder.py b/aider/coders/wholefile_coder.py index 840e8832413..d99e05dab25 100644 --- a/aider/coders/wholefile_coder.py +++ b/aider/coders/wholefile_coder.py @@ -38,14 +38,16 @@ def render_incremental_response(self, final): def update_files(self, mode="update"): content = self.partial_response_content - edited = set() chat_files = self.get_inchat_relative_files() output = [] lines = content.splitlines(keepends=True) + edits = [] + saw_fname = None fname = None + fname_source = None new_lines = [] for i, line in enumerate(lines): if line.startswith(self.fence[0]) or line.startswith(self.fence[1]): @@ -57,17 +59,17 @@ def update_files(self, mode="update"): if mode == "diff": output += self.do_live_diff(full_path, new_lines, True) - elif self.allowed_to_edit(fname): - edited.add(fname) - new_lines = "".join(new_lines) - self.io.write_text(full_path, new_lines) + else: + edits.append((fname, fname_source, new_lines)) fname = None + fname_source = None new_lines = [] continue # fname==None ... starting a new block if i > 0: + fname_source = "block" fname = lines[i - 1].strip() # Did gpt prepend a bogus dir? It especially likes to # include the path/to prefix from the one-shot example in @@ -77,8 +79,10 @@ def update_files(self, mode="update"): if not fname: # blank line? or ``` was on first line i==0 if saw_fname: fname = saw_fname + fname_source = "saw" elif len(chat_files) == 1: fname = chat_files[0] + fname_source = "chat" else: # TODO: sense which file it is by diff size raise ValueError( @@ -105,11 +109,22 @@ def update_files(self, mode="update"): return "\n".join(output) if fname: - full_path = self.allowed_to_edit(fname) - if full_path: - edited.add(fname) + edits.append((fname, fname_source, new_lines)) + + edited = set() + # process from most reliable filename, to least reliable + for source in ("block", "saw", "chat"): + for fname, fname_source, new_lines in edits: + if fname_source != source: + continue + # if a higher priority source already edited the file, skip + if fname in edited: + continue + + # we have a winner new_lines = "".join(new_lines) - self.io.write_text(full_path, new_lines) + if self.allowed_to_edit(fname, new_lines): + edited.add(fname) return edited diff --git a/tests/test_wholefile.py b/tests/test_wholefile.py index 83744c88508..d6ff75b4713 100644 --- a/tests/test_wholefile.py +++ b/tests/test_wholefile.py @@ -250,6 +250,38 @@ def test_update_files_earlier_filename(self): self.assertEqual(fname_a.read_text(), "after a\n") self.assertEqual(fname_b.read_text(), "after b\n") + def test_update_named_file_but_extra_unnamed_code_block(self): + sample_file = "hello.py" + new_content = "new\ncontent\ngoes\nhere\n" + + with open(sample_file, "w") as f: + f.write("Original content\n") + + # Initialize WholeFileCoder with the temporary directory + io = InputOutput(yes=True) + coder = WholeFileCoder(main_model=models.GPT35, io=io, fnames=[sample_file]) + + # Set the partial response content with the updated content + coder.partial_response_content = ( + f"Here's the modified `{sample_file}` file that implements the `accumulate`" + f" function as per the given instructions:\n\n```\n{new_content}```\n\nThis" + " implementation uses a list comprehension to apply the `operation` function to" + " each element of the `collection` and returns the resulting list.\n" + "Run it like this:\n\n" + "```\npython {sample_file}\n```\n\n" + ) + + # Call update_files method + edited_files = coder.update_files() + + # Check if the sample file was updated + self.assertIn(sample_file, edited_files) + + # Check if the content of the sample file was updated + with open(sample_file, "r") as f: + updated_content = f.read() + self.assertEqual(updated_content, new_content) + def test_full_edit(self): # Create a few temporary files _, file1 = tempfile.mkstemp()