From 8d2ccfcbd65b24376e5fb5a9b15e847751fc54f6 Mon Sep 17 00:00:00 2001
From: Nick Sullivan <nick@sullivanflock.com>
Date: Sun, 13 Aug 2023 19:17:40 +0200
Subject: [PATCH] Fix up apply patch/rebuild patch for new line issue.
 Parameterize the tests to make it easy to add more in the future

---
 aicodebot/patch.py                           |  10 +-
 aicodebot/prompts.py                         |  27 ++--
 pyproject.toml                               |   1 +
 tests/rebuild_patch_data/coder.patch         |  12 ++
 tests/rebuild_patch_data/coder.py            |   8 ++
 tests/rebuild_patch_data/coder_expected.py   |   8 ++
 tests/rebuild_patch_data/input.patch         |  12 ++
 tests/rebuild_patch_data/input.py            |  10 ++
 tests/rebuild_patch_data/input_expected.py   |  10 ++
 tests/rebuild_patch_data/prompts.patch       |  13 ++
 tests/rebuild_patch_data/prompts.py          |  11 ++
 tests/rebuild_patch_data/prompts_expected.py |  11 ++
 tests/test_input.py                          |   4 +-
 tests/test_patch.py                          | 133 ++++---------------
 14 files changed, 145 insertions(+), 125 deletions(-)
 create mode 100644 tests/rebuild_patch_data/coder.patch
 create mode 100644 tests/rebuild_patch_data/coder.py
 create mode 100644 tests/rebuild_patch_data/coder_expected.py
 create mode 100644 tests/rebuild_patch_data/input.patch
 create mode 100644 tests/rebuild_patch_data/input.py
 create mode 100644 tests/rebuild_patch_data/input_expected.py
 create mode 100644 tests/rebuild_patch_data/prompts.patch
 create mode 100644 tests/rebuild_patch_data/prompts.py
 create mode 100644 tests/rebuild_patch_data/prompts_expected.py

diff --git a/aicodebot/patch.py b/aicodebot/patch.py
index b7cd7cf..7fed197 100644
--- a/aicodebot/patch.py
+++ b/aicodebot/patch.py
@@ -17,7 +17,6 @@ def apply_patch(patch_string, is_rebuilt=False):
                     "apply",
                     "--verbose",
                     "--recount",
-                    "--inaccurate-eof",
                 ],
                 input=patch_string.encode("utf-8"),
                 check=True,
@@ -78,16 +77,16 @@ def rebuild_patch(patch_string):  # noqa: PLR0915
         # ------------------------- Parse the incoming patch ------------------------- #
         parsed_lines = []
         chunk_header = None
-        for line in patch_string.lstrip().splitlines():
+        for line in patch_string.splitlines():
             if chunk_header and not line.startswith(("+", "-", " ")):
                 # Sometimes the LM will add a context line without a space
                 # If we see that, we'll assume it's a context line
                 line = " " + line  # noqa: PLW2901
 
             parsed_line = Patch.parse_line(line)
+            if parsed_line.type == "chunk_header":
+                chunk_header = parsed_line.parsed
             parsed_lines.append(parsed_line)
-            if parsed_lines[-1].type == "chunk_header":
-                chunk_header = parsed_lines[-1].parsed
 
         # Check for critical fields
         source_file_line = next(line for line in parsed_lines if line.type == "source_file")
@@ -104,7 +103,7 @@ def rebuild_patch(patch_string):  # noqa: PLR0915
 
         start1 = chunk_header.start1
         first_change_line = next(line for line in parsed_lines if line.type in ("addition", "subtraction"))
-        lines_of_context = 3
+        lines_of_context = 1
 
         # ------------------------- Rebuild the context lines ------------------------ #
         # Get the correct start line from the first context line, by looking at the source file
@@ -118,6 +117,7 @@ def rebuild_patch(patch_string):  # noqa: PLR0915
             for i in range(start1 - 1, len(source_file_contents)):
                 if source_file_contents[i] == first_change_line.parsed:
                     first_change_line_number = i + 1
+                    start1 = first_change_line_number - lines_of_context
                     break
             else:
                 raise ValueError(f"Could not find first change line in source file: {first_change_line.parsed}")
diff --git a/aicodebot/prompts.py b/aicodebot/prompts.py
index c306dd3..ab790cb 100644
--- a/aicodebot/prompts.py
+++ b/aicodebot/prompts.py
@@ -200,10 +200,25 @@ def get_personality_prompt():
 
 PATCH_FORMAT_EXPLANATION = """
 To suggest a code change to the files in the local git repo, we use a unified diff format.
+The diff context is the output of the `git diff` command. It shows the changes that have been made.
+Lines starting with "-" are being removed. Lines starting with "+" are being added.
+Lines starting with " " (space) are unchanged. The file names are shown for context.
+
+ A line of code that is unchanged, that is being passed for context (starts with a space)
+ A second line of code that is unchanged, that is being passed for context (starts with a space)
+-A line of code that is being removed
++A line of code that is being added
+
+Before laying out the patch, write up a description of the change you want to make, to explain
+what you want to do.
 
 === Example ===
-Software Engineer: I want to add helpful header comments to the functions in file x.py
-AICodeBot: Ok, I've added helpful header comments to the functions in file x.py.
+Software Engineer: Fix the spelling mistake in x.py
+AICodeBot: Ok, I'll fix the spelling mistake in x.py
+
+Here's the change I am making:
+1. Remove the line "# Line with seplling mistake"
+2. Add the replacement line "# Line with spelling fixed"
 
 ```diff
 diff --git a/x.py b/x.py
@@ -212,15 +227,11 @@ def get_personality_prompt():
 @@ -1,3 +1,4 @@
 
 def foo():
-+    # New helpful header comment
+-    # Line with seplling mistake
++    # Line with spelling fixed
     pass
 ```
 === End Example ===
-
-In the above example, the engineer asked to add helpful header comments to the functions in file x.py.
-It was starting at line 1, and there were 3 lines of code in the original file. The changes spanned
-4 lines in the new file, so the chunk header was @@ -1,3 +1,4 @@. It's important to get the numbers
-in the chunk header correct, so that patch can be applied cleanly with git apply.
 """
 
 SIDEKICK_TEMPLATE = (
diff --git a/pyproject.toml b/pyproject.toml
index c86f60f..da8ff59 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,6 +84,7 @@ select = [
 target-version = "py311"
 # Certain errors we don't want to fix because they are too aggressive,
 # especially in the editor (removing variables we haven't used yet
+extend-exclude = ["tests/rebuild_patch_data"]
 unfixable = ["F401", "F841"]
 
 [tool.ruff.per-file-ignores]
diff --git a/tests/rebuild_patch_data/coder.patch b/tests/rebuild_patch_data/coder.patch
new file mode 100644
index 0000000..4fd27e9
--- /dev/null
+++ b/tests/rebuild_patch_data/coder.patch
@@ -0,0 +1,12 @@
+diff --git a/coder.py b/coder.py
+--- a/coder.py
++++ b/coder.py
+@@ -3,7 +3,7 @@
+ from pathlib import Path
+ from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
+ from types import SimpleNamespace
+-import fnmatch, mimetypes, re, subprocess, unidiff
++import fnmatch, mimetypes, re, subprocess
+
+
+ class Code
\ No newline at end of file
diff --git a/tests/rebuild_patch_data/coder.py b/tests/rebuild_patch_data/coder.py
new file mode 100644
index 0000000..ef57450
--- /dev/null
+++ b/tests/rebuild_patch_data/coder.py
@@ -0,0 +1,8 @@
+from aicodebot.helpers import exec_and_get_output, logger
+from aicodebot.lm import token_size
+from pathlib import Path
+from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
+from types import SimpleNamespace
+import fnmatch, mimetypes, re, subprocess, unidiff
+
+# Comment
diff --git a/tests/rebuild_patch_data/coder_expected.py b/tests/rebuild_patch_data/coder_expected.py
new file mode 100644
index 0000000..915ea2b
--- /dev/null
+++ b/tests/rebuild_patch_data/coder_expected.py
@@ -0,0 +1,8 @@
+from aicodebot.helpers import exec_and_get_output, logger
+from aicodebot.lm import token_size
+from pathlib import Path
+from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
+from types import SimpleNamespace
+import fnmatch, mimetypes, re, subprocess
+
+# Comment
diff --git a/tests/rebuild_patch_data/input.patch b/tests/rebuild_patch_data/input.patch
new file mode 100644
index 0000000..5d60135
--- /dev/null
+++ b/tests/rebuild_patch_data/input.patch
@@ -0,0 +1,12 @@
+diff --git a/input.py b/input.py
+--- a/input.py
++++ b/input.py
+@@ -1,7 +1,7 @@
+ from aicodebot.coder import Coder
+ from aicodebot.commands import commit, review
+-from aicodebot.lm import DEFAULT_RESPONSE_TOKENS, token_size
++from aicodebot.lm import token_size
+ from aicodebot.patch import Patch
+ from pathlib import Path
+ from prompt_toolkit import PromptSession
+ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
\ No newline at end of file
diff --git a/tests/rebuild_patch_data/input.py b/tests/rebuild_patch_data/input.py
new file mode 100644
index 0000000..523afc9
--- /dev/null
+++ b/tests/rebuild_patch_data/input.py
@@ -0,0 +1,10 @@
+from aicodebot.coder import Coder
+from aicodebot.commands import commit, review
+from aicodebot.lm import DEFAULT_RESPONSE_TOKENS, token_size
+from aicodebot.patch import Patch
+from pathlib import Path
+from prompt_toolkit import PromptSession
+from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.history import FileHistory
+from rich.panel import Panel
diff --git a/tests/rebuild_patch_data/input_expected.py b/tests/rebuild_patch_data/input_expected.py
new file mode 100644
index 0000000..83e6fab
--- /dev/null
+++ b/tests/rebuild_patch_data/input_expected.py
@@ -0,0 +1,10 @@
+from aicodebot.coder import Coder
+from aicodebot.commands import commit, review
+from aicodebot.lm import token_size
+from aicodebot.patch import Patch
+from pathlib import Path
+from prompt_toolkit import PromptSession
+from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.history import FileHistory
+from rich.panel import Panel
diff --git a/tests/rebuild_patch_data/prompts.patch b/tests/rebuild_patch_data/prompts.patch
new file mode 100644
index 0000000..8bd135f
--- /dev/null
+++ b/tests/rebuild_patch_data/prompts.patch
@@ -0,0 +1,13 @@
+diff --git a/prompts.py b/prompts.py
+--- a/prompts.py
++++ b/prompts.py
+@@ -6,7 +6,7 @@
+from langchain import PromptTemplate
+from langchain.output_parsers import PydanticOutputParser
+from pathlib import Path
+-from pydantic import BaseModel, Field
+-from types import SimpleNamespace
+-import arrow, functools, os, platform
++from pydantic import BaseModel, Field
++from types import SimpleNamespace
++import arrow, functools, os
\ No newline at end of file
diff --git a/tests/rebuild_patch_data/prompts.py b/tests/rebuild_patch_data/prompts.py
new file mode 100644
index 0000000..e5c69cb
--- /dev/null
+++ b/tests/rebuild_patch_data/prompts.py
@@ -0,0 +1,11 @@
+from aicodebot.coder import Coder
+from aicodebot.config import read_config
+from aicodebot.helpers import logger
+from langchain import PromptTemplate
+from langchain.output_parsers import PydanticOutputParser
+from pathlib import Path
+from pydantic import BaseModel, Field
+from types import SimpleNamespace
+import arrow, functools, os, platform
+
+# Comment
diff --git a/tests/rebuild_patch_data/prompts_expected.py b/tests/rebuild_patch_data/prompts_expected.py
new file mode 100644
index 0000000..7cb2b50
--- /dev/null
+++ b/tests/rebuild_patch_data/prompts_expected.py
@@ -0,0 +1,11 @@
+from aicodebot.coder import Coder
+from aicodebot.config import read_config
+from aicodebot.helpers import logger
+from langchain import PromptTemplate
+from langchain.output_parsers import PydanticOutputParser
+from pathlib import Path
+from pydantic import BaseModel, Field
+from types import SimpleNamespace
+import arrow, functools, os
+
+# Comment
diff --git a/tests/test_input.py b/tests/test_input.py
index 89f96ea..1b5cd5d 100644
--- a/tests/test_input.py
+++ b/tests/test_input.py
@@ -67,7 +67,7 @@ def test_apply_subcommand(chat, temp_git_repo):
     with in_temp_directory(temp_git_repo.working_dir):
         # Create a file to be modified
         mod_file = Path("mod_file.txt")
-        mod_file.write_text("AICodeBot is your coding sidekick.\nIt is here to make your coding life easier.")
+        mod_file.write_text("AICodeBot is your coding sidekick.\nIt is here to make your coding life easier.\n")
 
         # Create a patch to modify the file
         mod_patch = textwrap.dedent(
@@ -79,7 +79,7 @@ def test_apply_subcommand(chat, temp_git_repo):
              It is here to make your coding life easier.
             +It is now even better!
             """
-        )
+        ).lstrip()
         # Add the patch to the chat (simulating it coming in from the LM response)
         chat.diff_blocks = [mod_patch]
 
diff --git a/tests/test_patch.py b/tests/test_patch.py
index ddb303e..036db33 100644
--- a/tests/test_patch.py
+++ b/tests/test_patch.py
@@ -1,8 +1,7 @@
-from aicodebot.helpers import create_and_write_file
 from aicodebot.patch import Patch
 from pathlib import Path
 from tests.conftest import in_temp_directory
-import textwrap
+import pytest, shutil, textwrap
 
 
 def test_apply_patch(temp_git_repo):
@@ -13,7 +12,7 @@ def test_apply_patch(temp_git_repo):
 
         # Create a file to be modified
         mod_file = Path("mod_file.txt")
-        mod_file.write_text("AICodeBot is your coding sidekick.\nIt is here to make your coding life easier.")
+        mod_file.write_text("AICodeBot is your coding sidekick.\nIt is here to make your coding life easier.\n")
 
         # Create a file to be removed
         remove_file = Path("remove_file.txt")
@@ -58,118 +57,32 @@ def test_apply_patch(temp_git_repo):
         assert not remove_file.exists()
 
 
-def test_rebuild_patch(tmp_path):
-    # Use in_temp_directory for the test
-    with in_temp_directory(tmp_path):
-        # Set up the original file
-        Path(tmp_path / "aicodebot").mkdir()
-        create_and_write_file(
-            "aicodebot/prompts.py",
-            textwrap.dedent(
-                """
-                from aicodebot.coder import Coder
-                from aicodebot.config import read_config
-                from aicodebot.helpers import logger
-                from langchain import PromptTemplate
-                from langchain.output_parsers import PydanticOutputParser
-                from pathlib import Path
-                from pydantic import BaseModel, Field
-                from types import SimpleNamespace
-                import arrow, functools, os, platform
-
-
-                # Comment
-                """
-            ),
-        )
-        prompts_file = Path("aicodebot/prompts.py").read_text()
-        assert "platform" in prompts_file
-
-        # A few problems with this patch:
-        # 1. The chunk header is wrong (wrong line in the file and wrong number of lines)
-        # 2. No " " before the unchanged lines
-        # 3. Duplicated added/removed lines (that should just be unchanged)
-        # The original goal of the patch was to remove the "platform" import
-        bad_patch = textwrap.dedent(
-            """
-            diff --git a/aicodebot/prompts.py b/aicodebot/prompts.py
-            --- a/aicodebot/prompts.py
-            +++ b/aicodebot/prompts.py
-            @@ -6,7 +6,7 @@
-            from langchain import PromptTemplate
-            from langchain.output_parsers import PydanticOutputParser
-            from pathlib import Path
-            -from pydantic import BaseModel, Field
-            -from types import SimpleNamespace
-            -import arrow, functools, os, platform
-            +from pydantic import BaseModel, Field
-            +from types import SimpleNamespace
-            +import arrow, functools, os
-
-            """
-        )
-
-        print("Bad patch:\n", bad_patch)
-        rebuilt_patch = Patch.rebuild_patch(bad_patch)
-        print("Rebuilt patch:\n", rebuilt_patch)
-
-        # Apply the rebuilt patch
-        assert Patch.apply_patch(rebuilt_patch) is True
+@pytest.mark.parametrize(
+    "test_name, expected_chunk_header",
+    [
+        ("prompts", "@@ -6,5 +6,5 @@"),
+        ("coder", "@@ -5,3 +5,3 @@"),
+        ("input", "@@ -2,3 +2,3 @@"),
+    ],
+)
+def test_rebuild_patch_parameterized(tmp_path, test_name, expected_chunk_header):
+    test_files_dir = Path("tests/rebuild_patch_data")
+    shutil.copy(test_files_dir / f"{test_name}.py", tmp_path)
+    bad_patch = Path(test_files_dir / f"{test_name}.patch").read_text()
+    expected_result = Path(test_files_dir / f"{test_name}_expected.py").read_text()
 
-        assert "platform" not in Path("aicodebot/prompts.py").read_text()
-
-
-def test_rebuild_patch_coder(tmp_path):
     # Use in_temp_directory for the test
     with in_temp_directory(tmp_path):
-        # Set up the original file
-        file = "aicodebot/coder.py"
-        Path(tmp_path / "aicodebot").mkdir()
-        create_and_write_file(
-            file,
-            textwrap.dedent(
-                """
-                from aicodebot.helpers import exec_and_get_output, logger
-                from aicodebot.lm import token_size
-                from pathlib import Path
-                from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
-                from types import SimpleNamespace
-                import fnmatch, mimetypes, re, subprocess, unidiff
-
-
-                class Coder:
-                """
-            ).lstrip(),
-        )
-        assert "unidiff" in Path(file).read_text()
-
-        # A few problems with this patch:
-        # 1. The chunk header is wrong (wrong line in the file and wrong number of lines)
-        # 2. No " " before the unchanged lines
-        # 3. Duplicated added/removed lines (that should just be unchanged)
-        # The original goal of the patch was to remove the "platform" import
-        bad_patch = textwrap.dedent(
-            """
-            diff --git a/aicodebot/coder.py b/aicodebot/coder.py
-            --- a/aicodebot/coder.py
-            +++ b/aicodebot/coder.py
-            @@ -3,7 +3,7 @@
-             from pathlib import Path
-             from pygments.lexers import ClassNotFound, get_lexer_for_mimetype, guess_lexer_for_filename
-             from types import SimpleNamespace
-            -import fnmatch, mimetypes, re, subprocess, unidiff
-            +import fnmatch, mimetypes, re, subprocess
-
-
-             class Coder
-            """
-        ).lstrip()
-
-        print("Bad patch:\n", bad_patch)
+        print(f"Bad patch:\n{bad_patch}")
         rebuilt_patch = Patch.rebuild_patch(bad_patch)
-        print("Rebuilt patch:\n", rebuilt_patch)
+        print(f"Rebuilt patch:\n{rebuilt_patch}")
+
+        assert expected_chunk_header in rebuilt_patch
 
         # Apply the rebuilt patch
         assert Patch.apply_patch(rebuilt_patch) is True
 
-        assert "unidiff" not in Path(file).read_text()
+        patched_contents = Path(f"{test_name}.py").read_text()
+        print(f"Patched contents:\n{patched_contents}")
+        print(f"Expected result:\n{expected_result}")
+        assert patched_contents == expected_result