Improve patch rebuilding in Coder class

In this commit, we've added docstrings to several methods in the `Coder` class, `Chat` class, and `SidekickCompleter` class to improve code readability and maintainability. We've also enhanced the `rebuild_patch` method in the `Coder` class to handle missing chunk headers and rebuild context lines more accurately. This should make the patch application process more robust and reliable. 🛠️💡

aicodebot/coder.py

@@ -15,6 +15,7 @@ class Coder:
 
     @staticmethod
     def apply_patch(patch_string, is_rebuilt=False):
+        """Applies a patch to the local file system using git apply."""
         try:
             result = subprocess.run(
                 [
@@ -324,6 +325,7 @@ def rebuild_patch(patch_string):  # noqa: PLR0915
         This function looks at the intent of the patch and rebuilds it in a [hopefully] correct format."""
 
         def parse_line(line):  # noqa: PLR0911
+            """Parse a line of the patch and return a SimpleNamespace with the line, type, and parsed line."""
             if line.startswith(("diff --git", "index")):
                 return SimpleNamespace(line=line, type="header", parsed=line)
             elif line.startswith("---"):
@@ -352,6 +354,7 @@ def parse_line(line):  # noqa: PLR0911
             else:
                 raise ValueError(f"Invalid line: '{line}'")
 
+        # ------------------------- Parse the incoming patch ------------------------- #
         parsed_lines = []
         chunk_header = None
         for line in patch_string.lstrip().splitlines():
@@ -361,11 +364,11 @@ def parse_line(line):  # noqa: PLR0911
                 line = " " + line  # noqa: PLW2901
 
             parsed_line = parse_line(line)
-            logger.error(f"{parsed_line.type}: {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")
         if not source_file_line:
             raise ValueError("No source file found in patch")
@@ -375,20 +378,48 @@ def parse_line(line):  # noqa: PLR0911
             raise ValueError("No context line found in patch")
 
         if not chunk_header:
-            # This shouldn't happen, but we should be able to recover
+            # Chunk header missing. This shouldn't happen, but we should be able to recover
             chunk_header = SimpleNamespace(start1=0, count1=0, start2=0, count2=0)
+
         start1 = chunk_header.start1
+        first_change_line = next(line for line in parsed_lines if line.type in ("addition", "subtraction"))
+        lines_of_context = 3
 
+        # ------------------------- Rebuild the context lines ------------------------ #
         # Get the correct start line from the first context line, by looking at the source file
         source_file = source_file_line.parsed
+        source_file_contents = []
         if source_file != "/dev/null" and Path(source_file).exists():
             source_file_contents = Path(source_file).read_text().splitlines()
-            for i in range(0, len(source_file_contents)):
-                if source_file_contents[i] == first_context_line.parsed:
-                    start1 = i + 1
+
+            # Determine the correct line of the first change
+            # We will start looking at start1 - 1, and walk until we find it
+            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
                     break
+            else:
+                raise ValueError(f"Could not find first change line in source file: {first_change_line.parsed}")
+
+            # Disregard the existing context lines from the parsed lines
+            parsed_lines = [line for line in parsed_lines if line.type != "context"]
+
+            # Add x lines of context before the first change
+            for i in range(first_change_line_number - lines_of_context, first_change_line_number):
+                # Get the index number of the first changed line in parsed_lines
+                first_change_line_index = next(
+                    i for i, line in enumerate(parsed_lines) if line.type in ("addition", "subtraction")
+                )
+                parsed_lines.insert(first_change_line_index, parse_line(f" {source_file_contents[i-1]}"))
+
+            # Add x lines of context after the last change
+            number_of_subtractions = len([line for line in parsed_lines if line.type == "subtraction"])
+            start_trailing_context = first_change_line_number + number_of_subtractions
+            for i in range(start_trailing_context, start_trailing_context + lines_of_context):
+                parsed_lines.append(parse_line(f" {source_file_contents[i-1]}"))
+
+        # ------------------------- Rebuild the chunk header ------------------------- #
 
-        # Calculate the new chunk header
         # Add up the number of context lines, additions, and subtractions
         # This will be the new count1 and count2
         start2 = start1
@@ -401,7 +432,8 @@ def parse_line(line):  # noqa: PLR0911
 
         new_chunk_header = f"@@ -{start1},{count1} +{start2},{count2} @@"
 
-        # Rebuild the patch
+        # ----------------------------- Rebuild the patch ---------------------------- #
+
         new_patch = []
         for line in parsed_lines:
             if line.type == "chunk_header":

aicodebot/input.py

@@ -21,6 +21,7 @@ def __init__(self, console, files):
         self.files = set(files)
 
     def parse_human_input(self, human_input):  # noqa: PLR0911, PLR0915
+        """Parse the human input and handle any special commands."""
         human_input = human_input.strip()
 
         if not human_input:
@@ -172,6 +173,7 @@ def project_files(self):
         return self._project_files
 
     def get_completions(self, document, complete_event):
+        """yield prompt_toolkit Completion objects for the current input"""
         # Get the text before the cursor
         text = document.text_before_cursor
 

tests/test_coder.py

@@ -257,6 +257,8 @@ def test_rebuild_patch(tmp_path):
                 from types import SimpleNamespace
                 import arrow, functools, os, platform
 
+
+                # Comment
                 """
             ),
         )
@@ -287,17 +289,17 @@ def test_rebuild_patch(tmp_path):
             """
         )
 
+        print("Bad patch:\n", bad_patch)
         rebuilt_patch = Coder.rebuild_patch(bad_patch)
+        print("Rebuilt patch:\n", rebuilt_patch)
 
         # Apply the rebuilt patch
-        assert Coder.apply_patch(rebuilt_patch)
+        assert Coder.apply_patch(rebuilt_patch) is True
 
         assert "platform" not in Path("aicodebot/prompts.py").read_text()
 
 
 def test_rebuild_patch_coder(tmp_path):
-    return  # This test fails, checking it in for future use case
-
     # Use in_temp_directory for the test
     with in_temp_directory(tmp_path):
         # Set up the original file
@@ -343,9 +345,11 @@ class Coder
             """
         ).lstrip()
 
+        print("Bad patch:\n", bad_patch)
         rebuilt_patch = Coder.rebuild_patch(bad_patch)
+        print("Rebuilt patch:\n", rebuilt_patch)
 
         # Apply the rebuilt patch
-        assert Coder.apply_patch(rebuilt_patch)
+        assert Coder.apply_patch(rebuilt_patch) is True
 
         assert "unidiff" not in Path(file).read_text()