diff --git a/pyproject.toml b/pyproject.toml index 4c7fd94a..b36da8ed 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/codemodder/dependency_management/base_dependency_writer.py b/src/codemodder/dependency_management/base_dependency_writer.py index 22236984..07129b5d 100644 --- a/src/codemodder/dependency_management/base_dependency_writer.py +++ b/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 = [] diff --git a/src/codemodder/dependency_management/dependency_manager.py b/src/codemodder/dependency_management/dependency_manager.py index d626e48a..73d40238 100644 --- a/src/codemodder/dependency_management/dependency_manager.py +++ b/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 @@ def write( 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 diff --git a/src/codemodder/dependency_management/pyproject_writer.py b/src/codemodder/dependency_management/pyproject_writer.py new file mode 100644 index 00000000..e869fe28 --- /dev/null +++ b/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) diff --git a/src/codemodder/dependency_management/requirements_txt_writer.py b/src/codemodder/dependency_management/requirements_txt_writer.py index 1c59c278..aaa2c884 100644 --- a/src/codemodder/dependency_management/requirements_txt_writer.py +++ b/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, diff --git a/src/codemodder/diff.py b/src/codemodder/diff.py index ec452d59..d0aa25d8 100644 --- a/src/codemodder/diff.py +++ b/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 = [ diff --git a/tests/dependency_management/test_pyproject_writer.py b/tests/dependency_management/test_pyproject_writer.py new file mode 100644 index 00000000..fa8fe5b3 --- /dev/null +++ b/tests/dependency_management/test_pyproject_writer.py @@ -0,0 +1,209 @@ +from textwrap import dedent +import pytest + +from codemodder.dependency_management.pyproject_writer import PyprojectWriter +from codemodder.dependency import DefusedXML, Security +from codemodder.project_analysis.file_parsers.package_store import PackageStore + + +@pytest.mark.parametrize("dry_run", [True, False]) +def test_update_pyproject_dependencies(tmpdir, dry_run): + orig_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + + [project] + dynamic = ["version"] + name = "codemodder" + requires-python = ">=3.10.0" + readme = "README.md" + dependencies = [ + "libcst~=1.1.0", + "pylint~=3.0.0", + "PyYAML~=6.0.0", + ] + """ + + pyproject_toml = tmpdir.join("pyproject.toml") + pyproject_toml.write(dedent(orig_pyproject)) + + store = PackageStore( + type="requirements.txt", + file=str(pyproject_toml), + dependencies=set(), + py_versions=[">=3.10.0"], + ) + + writer = PyprojectWriter(store, tmpdir) + dependencies = [DefusedXML, Security] + changeset = writer.write(dependencies, dry_run=dry_run) + + updated_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + + [project] + dynamic = ["version"] + name = "codemodder" + requires-python = ">=3.10.0" + readme = "README.md" + dependencies = [ + "libcst~=1.1.0", + "pylint~=3.0.0", + "PyYAML~=6.0.0", + "defusedxml~=0.7.1", + "security~=1.2.0", + ] + """ + + assert pyproject_toml.read() == ( + dedent(orig_pyproject) if dry_run else dedent(updated_pyproject) + ) + + assert changeset is not None + assert changeset.path == pyproject_toml.basename + res = ( + "--- \n" + "+++ \n" + "@@ -11,5 +11,7 @@\n" + """ "libcst~=1.1.0",\n""" + """ "pylint~=3.0.0",\n""" + """ "PyYAML~=6.0.0",\n""" + """+ "defusedxml~=0.7.1",\n""" + """+ "security~=1.2.0",\n""" + " ]\n " + ) + assert changeset.diff == res + assert len(changeset.changes) == 2 + change_one = changeset.changes[0] + + assert change_one.lineNumber == 14 + assert change_one.description == DefusedXML.build_description() + assert change_one.properties == { + "contextual_description": True, + "contextual_description_position": "right", + } + change_two = changeset.changes[1] + assert change_two.lineNumber == 15 + assert change_two.description == Security.build_description() + assert change_two.properties == { + "contextual_description": True, + "contextual_description_position": "right", + } + + +def test_add_same_dependency_only_once(tmpdir): + orig_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + + [project] + dynamic = ["version"] + name = "codemodder" + requires-python = ">=3.10.0" + readme = "README.md" + dependencies = [ + "libcst~=1.1.0", + "pylint~=3.0.0", + "PyYAML~=6.0.0", + ] + """ + + pyproject_toml = tmpdir.join("pyproject.toml") + pyproject_toml.write(dedent(orig_pyproject)) + + store = PackageStore( + type="requirements.txt", + file=str(pyproject_toml), + dependencies=set(), + py_versions=[">=3.10.0"], + ) + + writer = PyprojectWriter(store, tmpdir) + dependencies = [Security, Security] + writer.write(dependencies) + + updated_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + + [project] + dynamic = ["version"] + name = "codemodder" + requires-python = ">=3.10.0" + readme = "README.md" + dependencies = [ + "libcst~=1.1.0", + "pylint~=3.0.0", + "PyYAML~=6.0.0", + "security~=1.2.0", + ] + """ + + assert pyproject_toml.read() == dedent(updated_pyproject) + + +def test_dont_add_existing_dependency(tmpdir): + orig_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + + [project] + dynamic = ["version"] + name = "codemodder" + requires-python = ">=3.10.0" + readme = "README.md" + dependencies = [ + "libcst~=1.1.0", + "pylint~=3.0.0", + "PyYAML~=6.0.0", + ] + """ + + pyproject_toml = tmpdir.join("pyproject.toml") + pyproject_toml.write(dedent(orig_pyproject)) + + store = PackageStore( + type="requirements.txt", + file=str(pyproject_toml), + dependencies=set([Security.requirement]), + py_versions=[">=3.10.0"], + ) + + writer = PyprojectWriter(store, tmpdir) + dependencies = [Security] + writer.write(dependencies) + + assert pyproject_toml.read() == dedent(orig_pyproject) + + +def test_pyproject_no_dependencies(tmpdir): + orig_pyproject = """\ + [build-system] + requires = ["setuptools", "setuptools_scm>=8"] + build-backend = "setuptools.build_meta" + [project] + name = "codemodder" + """ + + pyproject_toml = tmpdir.join("pyproject.toml") + pyproject_toml.write(dedent(orig_pyproject)) + + store = PackageStore( + type="requirements.txt", + file=str(pyproject_toml), + dependencies=set(), + py_versions=[">=3.10.0"], + ) + + writer = PyprojectWriter(store, tmpdir) + dependencies = [Security] + + writer.write(dependencies) + + assert pyproject_toml.read() == dedent(orig_pyproject) diff --git a/tests/dependency_management/test_requirements_txt_writer.py b/tests/dependency_management/test_requirements_txt_writer.py index 9a376e76..fd5eb541 100644 --- a/tests/dependency_management/test_requirements_txt_writer.py +++ b/tests/dependency_management/test_requirements_txt_writer.py @@ -7,11 +7,6 @@ from codemodder.dependency import DefusedXML, Security -@pytest.fixture(autouse=True, scope="module") -def disable_write_dependencies(): - """Override fixture from conftest.py in order to allow testing""" - - class TestRequirementsTxtWriter: @pytest.mark.parametrize("dry_run", [True, False]) def test_add_dependencies_preserve_comments(self, tmpdir, dry_run):