From cded591f6fe1e484127818364ab015abab9e9a1e Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Fri, 17 Nov 2023 12:32:16 -0500
Subject: [PATCH] Handle requirements.txt files without trailing newlines

---
 src/codemodder/codemodder.py         | 19 +++++++------------
 src/codemodder/dependency_manager.py | 20 +++++++++++---------
 src/codemodder/diff.py               | 12 ++++++++++++
 tests/test_dependency_manager.py     | 17 ++++++++++++++---
 4 files changed, 44 insertions(+), 24 deletions(-)
 create mode 100644 src/codemodder/diff.py

diff --git a/src/codemodder/codemodder.py b/src/codemodder/codemodder.py
index a517c8ca..c09ff38a 100644
--- a/src/codemodder/codemodder.py
+++ b/src/codemodder/codemodder.py
@@ -1,6 +1,5 @@
 from concurrent.futures import ThreadPoolExecutor
 import datetime
-import difflib
 import itertools
 import logging
 import os
@@ -17,6 +16,7 @@
 from codemodder.change import ChangeSet
 from codemodder.code_directory import file_line_patterns, match_files
 from codemodder.context import CodemodExecutionContext
+from codemodder.diff import create_diff as create_diff_from_lines
 from codemodder.executor import CodemodExecutorWrapper
 from codemodder.project_analysis.python_repo_manager import PythonRepoManager
 from codemodder.report.codetf_reporter import report_default
@@ -49,18 +49,13 @@ def find_semgrep_results(
 
 
 def create_diff(original_tree: cst.Module, new_tree: cst.Module) -> str:
-    diff_lines = list(
-        difflib.unified_diff(
-            original_tree.code.splitlines(keepends=True),
-            new_tree.code.splitlines(keepends=True),
-        )
+    """
+    Create a diff between the original and output trees.
+    """
+    return create_diff_from_lines(
+        original_tree.code.splitlines(keepends=True),
+        new_tree.code.splitlines(keepends=True),
     )
-    # All but the last diff line should end with a newline
-    # The last diff line should be preserved as-is (with or without a newline)
-    diff_lines = [
-        line if line.endswith("\n") else line + "\n" for line in diff_lines[:-1]
-    ] + [diff_lines[-1]]
-    return "".join(diff_lines)
 
 
 def apply_codemod_to_file(
diff --git a/src/codemodder/dependency_manager.py b/src/codemodder/dependency_manager.py
index 6398bf8c..bdd75d9f 100644
--- a/src/codemodder/dependency_manager.py
+++ b/src/codemodder/dependency_manager.py
@@ -2,10 +2,10 @@
 from pathlib import Path
 from typing import Optional
 
-import difflib
 from packaging.requirements import Requirement
 
 from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
+from codemodder.diff import create_diff
 from codemodder.dependency import Dependency
 
 
@@ -38,13 +38,18 @@ def write(self, dry_run: bool = False) -> Optional[ChangeSet]:
         if not (self.dependency_file and self._new_requirements):
             return None
 
-        updated = self._lines + self.new_requirements + ["\n"]
+        original_lines = self._lines.copy()
+        if not original_lines[-1].endswith("\n"):
+            original_lines[-1] += "\n"
 
-        diff = "".join(difflib.unified_diff(self._lines, updated))
+        requirement_lines = [f"{req}\n" for req in self.new_requirements]
+
+        updated = original_lines + requirement_lines
+        diff = create_diff(self._lines, updated)
 
         changes = [
             Change(
-                lineNumber=len(self._lines) + i + 1,
+                lineNumber=len(original_lines) + i + 1,
                 description=dep.build_description(),
                 properties={"contextual_description": True},
                 packageActions=[
@@ -56,11 +61,8 @@ def write(self, dry_run: bool = False) -> Optional[ChangeSet]:
 
         if not dry_run:
             with open(self.dependency_file, "w", encoding="utf-8") as f:
-                f.writelines(self._lines)
-                if not self._lines[-1].endswith("\n"):
-                    f.write("\n")
-
-                f.writelines([f"{line}\n" for line in self.new_requirements])
+                f.writelines(original_lines)
+                f.writelines(requirement_lines)
 
         self.dependency_file_changed = True
         return ChangeSet(
diff --git a/src/codemodder/diff.py b/src/codemodder/diff.py
new file mode 100644
index 00000000..ec452d59
--- /dev/null
+++ b/src/codemodder/diff.py
@@ -0,0 +1,12 @@
+import difflib
+
+
+def create_diff(original_lines: list[str], new_lines: list[str]) -> str:
+    diff_lines = list(difflib.unified_diff(original_lines, new_lines))
+
+    # All but the last diff line should end with a newline
+    # The last diff line should be preserved as-is (with or without a newline)
+    diff_lines = [
+        line if line.endswith("\n") else line + "\n" for line in diff_lines[:-1]
+    ] + [diff_lines[-1]]
+    return "".join(diff_lines)
diff --git a/tests/test_dependency_manager.py b/tests/test_dependency_manager.py
index 2700b584..14017d35 100644
--- a/tests/test_dependency_manager.py
+++ b/tests/test_dependency_manager.py
@@ -40,11 +40,11 @@ def test_add_dependency_preserve_comments(self, tmpdir, dry_run):
         assert changeset.diff == (
             "--- \n"
             "+++ \n"
-            "@@ -1,3 +1,5 @@\n"
+            "@@ -1,3 +1,4 @@\n"
             " # comment\n"
             " \n"
             " requests\n"
-            "+defusedxml~=0.7.1+\n"
+            "+defusedxml~=0.7.1\n"
         )
         assert len(changeset.changes) == 1
         assert changeset.changes[0].lineNumber == 4
@@ -96,7 +96,18 @@ def test_dependency_file_no_terminating_newline(self, tmpdir):
 
         dm = DependencyManager(Path(tmpdir))
         dm.add([Security])
-        dm.write()
+        changeset = dm.write()
+
+        assert changeset is not None
+        assert changeset.diff == (
+            "--- \n"
+            "+++ \n"
+            "@@ -1,2 +1,3 @@\n"
+            " requests\n"
+            "-whatever\n"
+            "+whatever\n"
+            "+security~=1.2.0\n"
+        )
 
         assert dependency_file.read_text(encoding="utf-8") == (
             "requests\nwhatever\nsecurity~=1.2.0\n"