From 737f08036764a10a7c9341564977c4c4598c493f Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 13 Jan 2025 09:24:12 +0000
Subject: [PATCH] refactor: improve diff generation in FixIssue and ModifyCode
 steps

Co-Authored-By: Patched <tech@patched.codes>
---
 patchwork/steps/FixIssue/FixIssue.py     | 11 +++---
 patchwork/steps/ModifyCode/ModifyCode.py | 43 +++++++++++++++++-------
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/patchwork/steps/FixIssue/FixIssue.py b/patchwork/steps/FixIssue/FixIssue.py
index c9a855d9..efca7e43 100644
--- a/patchwork/steps/FixIssue/FixIssue.py
+++ b/patchwork/steps/FixIssue/FixIssue.py
@@ -127,18 +127,17 @@ def run(self):
             if isinstance(tool, CodeEditTool):
                 cwd = Path.cwd()
                 modified_files = [file_path.relative_to(cwd) for file_path in tool.tool_records["modified_files"]]
-                # Get the diff for each modified file
+                # Get the diff for each modified file using git
                 modified_files_with_diffs = []
+                repo = Repo(cwd, search_parent_directories=True)
                 for file in modified_files:
-                    # Read the current content of the file
                     file_path = Path(file)
                     if file_path.exists():
-                        current_content = file_path.read_text()
-                        # For now, we'll use the entire content as the diff since we don't have the original
-                        # In a future enhancement, we could store the original content and generate a proper diff
+                        # Get the diff using git
+                        diff = repo.git.diff('HEAD', str(file))
                         modified_files_with_diffs.append({
                             "path": str(file),
-                            "diff": current_content
+                            "diff": diff
                         })
                 return dict(modified_files=modified_files_with_diffs)
         return dict()
diff --git a/patchwork/steps/ModifyCode/ModifyCode.py b/patchwork/steps/ModifyCode/ModifyCode.py
index 3a302d28..894e29e0 100644
--- a/patchwork/steps/ModifyCode/ModifyCode.py
+++ b/patchwork/steps/ModifyCode/ModifyCode.py
@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import difflib
+import shutil
+import tempfile
 from pathlib import Path
 
 from patchwork.step import Step, StepStatus
@@ -89,20 +92,34 @@ def run(self) -> dict:
             if new_code is None:
                 continue
 
-            # Store the original content before replacement
-            original_content = ""
-            if Path(uri).exists():
-                with open(uri, 'r') as f:
-                    original_lines = f.readlines()
-                    if start_line is not None and end_line is not None:
-                        original_content = ''.join(original_lines[start_line:end_line])
-                    else:
-                        original_content = ''.join(original_lines)
-
-            replace_code_in_file(uri, start_line, end_line, new_code)
+            # Get the original content for diffing
+            file_path = Path(uri)
+            original_path = None
+            diff = ""
             
-            # Create a basic diff format showing the changes
-            diff = f"--- Original\n+++ Modified\n-{original_content}\n+{new_code}"
+            if file_path.exists():
+                # Create a temporary copy of the original file
+                with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp_file:
+                    shutil.copy2(uri, tmp_file.name)
+                    original_path = tmp_file.name
+
+                # Apply the changes
+                replace_code_in_file(uri, start_line, end_line, new_code)
+                
+                # Generate a proper unified diff
+                with open(original_path, 'r') as f1, open(uri, 'r') as f2:
+                    diff = ''.join(difflib.unified_diff(
+                        f1.readlines(),
+                        f2.readlines(),
+                        fromfile='a/' + str(file_path),
+                        tofile='b/' + str(file_path)
+                    ))
+                
+                # Clean up temporary file
+                Path(original_path).unlink()
+            else:
+                # If file doesn't exist, just store the new code as the diff
+                diff = f"+++ {file_path}\n{new_code}"
             
             modified_code_file = dict(
                 path=uri,