Patched patchwork/steps/ModifyCode/ModifyCode.py

Author: patched.codes[bot]

patchwork/steps/ModifyCode/ModifyCode.py

@@ -7,7 +7,7 @@
 
 def save_file_contents(file_path, content):
     """Utility function to save content to a file."""
-    with open(file_path, "w") as file:
+    with open(file_path, "w", newline='') as file:
         file.write(content)
 
 
@@ -32,30 +32,75 @@ def handle_indent(src: list[str], target: list[str], start: int, end: int) -> li
     return [indent + line for line in target]
 
 
+def detect_line_ending(content: bytes) -> str:
+    """
+    Detect the dominant line ending in a file.
+    Returns '\r\n' for CRLF or '\n' for LF.
+    """
+    if b'\r\n' in content:
+        return '\r\n'
+    return '\n'
+
+def normalize_line_endings(text: str, line_ending: str) -> str:
+    """Normalize all line endings in text to the specified ending."""
+    # First convert all line endings to \n
+    text = text.replace('\r\n', '\n')
+    # Then convert all \n to desired line ending
+    if line_ending != '\n':
+        text = text.replace('\n', line_ending)
+    return text
+
 def replace_code_in_file(
     file_path: str,
     start_line: int | None,
     end_line: int | None,
     new_code: str,
 ) -> None:
     path = Path(file_path)
-    new_code_lines = new_code.splitlines(keepends=True)
-    if len(new_code_lines) > 0 and not new_code_lines[-1].endswith("\n"):
-        new_code_lines[-1] += "\n"
-
+    
+    # Detect line endings from existing file
+    if path.exists():
+        with open(file_path, 'rb') as f:
+            content = f.read()
+            line_ending = detect_line_ending(content)
+            original_content = content.decode('utf-8')
+    else:
+        line_ending = '\n'  # Default for new files
+        original_content = ""
+    
+    # Normalize the new code to match the file's line endings
+    normalized_new_code = normalize_line_endings(new_code, line_ending)
+    
     if path.exists() and start_line is not None and end_line is not None:
-        """Replaces specified lines in a file with new code."""
-        text = path.read_text()
-
-        lines = text.splitlines(keepends=True)
-
-        # Insert the new code at the start line after converting it into a list of lines
-        lines[start_line:end_line] = handle_indent(lines, new_code_lines, start_line, end_line)
+        # Split preserving line endings
+        lines = original_content.splitlines(keepends=True)
+        new_lines = normalized_new_code.splitlines(keepends=True)
+        
+        # Handle indentation
+        indented_new_code = handle_indent(lines, new_lines, start_line, end_line)
+        
+        # Ensure all new lines have the correct ending
+        final_new_code = []
+        for line in indented_new_code:
+            if not line.endswith(line_ending):
+                if line.endswith('\n') or line.endswith('\r\n'):
+                    line = line.rstrip('\r\n') + line_ending
+                else:
+                    line = line + line_ending
+            final_new_code.append(line)
+        
+        # Replace the lines
+        lines[start_line:end_line] = final_new_code
+        final_content = ''.join(lines)
     else:
-        lines = new_code_lines
-
-    # Save the modified contents back to the file
-    save_file_contents(file_path, "".join(lines))
+        new_lines = normalized_new_code.splitlines(keepends=True)
+        if not new_lines[-1].endswith(line_ending):
+            new_lines[-1] = new_lines[-1].rstrip('\r\n') + line_ending
+        final_content = ''.join(new_lines)
+    
+    # Write the content back
+    with open(file_path, 'wb') as f:
+        f.write(final_content.encode('utf-8'))
 
 
 class ModifyCode(Step):