pyproject.toml writer

pyproject.toml

@@ -19,6 +19,7 @@ dependencies = [
     "PyYAML~=6.0.0",
     "semgrep~=1.50.0",
     "toml~=0.10.2",
+    "tomlkit~=0.12.0",
     "wrapt~=1.16.0",
 ]
 keywords = ["codemod", "codemods", "security", "fix", "fixes"]

src/codemodder/dependency_management/base_dependency_writer.py

@@ -16,11 +16,19 @@ def __init__(self, dependency_store: PackageStore, parent_directory: Path):
         self.parent_directory = parent_directory
 
     @abstractmethod
-    def write(
+    def add_to_file(
         self, dependencies: list[Dependency], dry_run: bool = False
     ) -> Optional[ChangeSet]:
         pass
 
+    def write(
+        self, dependencies: list[Dependency], dry_run: bool = False
+    ) -> Optional[ChangeSet]:
+        new_dependencies = self.add(dependencies)
+        if new_dependencies:
+            return self.add_to_file(new_dependencies, dry_run)
+        return None
+
     def add(self, dependencies: list[Dependency]) -> list[Dependency]:
         """add any number of dependencies to the end of list of dependencies."""
         new = []

src/codemodder/dependency_management/dependency_manager.py

@@ -4,6 +4,7 @@
 from codemodder.dependency_management.requirements_txt_writer import (
     RequirementsTxtWriter,
 )
+from codemodder.dependency_management.pyproject_writer import PyprojectWriter
 from codemodder.project_analysis.file_parsers.package_store import PackageStore
 from pathlib import Path
 
@@ -27,6 +28,10 @@
                 return RequirementsTxtWriter(
                     self.dependencies_store, self.parent_directory
                 ).write(dependencies, dry_run)
+            case "pyproject.toml":
+                return PyprojectWriter(
+                    self.dependencies_store, self.parent_directory
+                ).write(dependencies, dry_run)
             case "setup.py":
                 pass
         return None

src/codemodder/dependency_management/pyproject_writer.py

@@ -0,0 +1,55 @@
+import tomlkit
+from typing import Optional
+from copy import deepcopy
+from codemodder.dependency import Dependency
+from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
+from codemodder.dependency_management.base_dependency_writer import DependencyWriter
+from codemodder.diff import create_diff_and_linenums
+
+
+class PyprojectWriter(DependencyWriter):
+    def add_to_file(
+        self, dependencies: list[Dependency], dry_run: bool = False
+    ) -> Optional[ChangeSet]:
+        pyproject = self._parse_file()
+        original = deepcopy(pyproject)
+
+        try:
+            pyproject["project"]["dependencies"].extend(
+                [f"{dep.requirement}" for dep in dependencies]
+            )
+        except tomlkit.exceptions.NonExistentKey:
+            return None
+
+        diff, added_line_nums = create_diff_and_linenums(
+            tomlkit.dumps(original).split("\n"), tomlkit.dumps(pyproject).split("\n")
+        )
+
+        if not dry_run:
+            with open(self.path, "w", encoding="utf-8") as f:
+                tomlkit.dump(pyproject, f)
+
+        changes = [
+            Change(
+                lineNumber=added_line_nums[i],
+                description=dep.build_description(),
+                # Contextual comments should be added to the right side of split diffs
+                properties={
+                    "contextual_description": True,
+                    "contextual_description_position": "right",
+                },
+                packageActions=[
+                    PackageAction(Action.ADD, Result.COMPLETED, str(dep.requirement))
+                ],
+            )
+            for i, dep in enumerate(dependencies)
+        ]
+        return ChangeSet(
+            str(self.path.relative_to(self.parent_directory)),
+            diff,
+            changes=changes,
+        )
+
+    def _parse_file(self):
+        with open(self.path, encoding="utf-8") as f:
+            return tomlkit.load(f)

src/codemodder/dependency_management/requirements_txt_writer.py

@@ -3,19 +3,12 @@
 from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
 from codemodder.diff import create_diff
 from codemodder.dependency import Dependency
-from packaging.requirements import Requirement
 
 
 class RequirementsTxtWriter(DependencyWriter):
-    def write(
+    def add_to_file(
         self, dependencies: list[Dependency], dry_run: bool = False
     ) -> Optional[ChangeSet]:
-        new_dependencies = self.add(dependencies)
-        if new_dependencies:
-            return self.add_to_file(new_dependencies, dry_run)
-        return None
-
-    def add_to_file(self, dependencies: list[Requirement], dry_run: bool):
         lines = self._parse_file()
         original_lines = lines.copy()
         if not original_lines[-1].endswith("\n"):
@@ -26,6 +19,10 @@ def add_to_file(self, dependencies: list[Requirement], dry_run: bool):
 
         diff = create_diff(original_lines, updated_lines)
 
+        if not dry_run:
+            with open(self.path, "w", encoding="utf-8") as f:
+                f.writelines(updated_lines)
+
         changes = [
             Change(
                 lineNumber=len(original_lines) + i + 1,
@@ -41,11 +38,6 @@ def add_to_file(self, dependencies: list[Requirement], dry_run: bool):
             )
             for i, dep in enumerate(dependencies)
         ]
-
-        if not dry_run:
-            with open(self.path, "w", encoding="utf-8") as f:
-                f.writelines(updated_lines)
-
         return ChangeSet(
             str(self.path.relative_to(self.parent_directory)),
             diff,

src/codemodder/diff.py

@@ -3,7 +3,48 @@
 
 def create_diff(original_lines: list[str], new_lines: list[str]) -> str:
     diff_lines = list(difflib.unified_diff(original_lines, new_lines))
+    return difflines_to_str(diff_lines)
 
+
+def create_diff_and_linenums(
+    original_lines: list[str], new_lines: list[str]
+) -> tuple[str, list[int]]:
+    diff_lines = list(difflib.unified_diff(original_lines, new_lines))
+    return difflines_to_str(diff_lines), calc_new_line_nums(diff_lines)
+
+
+def calc_new_line_nums(diff_lines: list[str]) -> list[int]:
+    if not diff_lines:
+        return []
+
+    added_line_nums = []
+    current_line_number = 0
+
+    for line in diff_lines:
+        if line.startswith("@@"):
+            # Extract the starting line number for the updated file from the diff metadata.
+            # The format is @@ -x,y +a,b @@, where a is the starting line number in the updated file.
+            start_line = line.split(" ")[2]
+            current_line_number = (
+                int(start_line.split(",")[0][1:]) - 1
+            )  # Subtract 1 because line numbers are 1-indexed
+
+        elif line.startswith("+"):
+            # Increment line number for each line in the updated file
+            current_line_number += 1
+            if not line.startswith("++"):  # Ignore the diff metadata lines
+                added_line_nums.append(current_line_number)
+
+        elif not line.startswith("-"):
+            # Increment line number for unchanged/context lines
+            current_line_number += 1
+
+    return added_line_nums
+
+
+def difflines_to_str(diff_lines: list[str]) -> str:
+    if not diff_lines:
+        return ""
     # 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 = [