Setup py writer

src/codemodder/codemodder.py

@@ -16,7 +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.diff import create_diff_from_tree
 from codemodder.executor import CodemodExecutorWrapper
 from codemodder.project_analysis.python_repo_manager import PythonRepoManager
 from codemodder.report.codetf_reporter import report_default
@@ -48,16 +48,6 @@ def find_semgrep_results(
     return run_semgrep(context, yaml_files)
 
 
-def create_diff(original_tree: cst.Module, new_tree: cst.Module) -> str:
-    """
-    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),
-    )
-
-
 def apply_codemod_to_file(
     base_directory: Path,
     file_context,
@@ -78,7 +68,7 @@ def apply_codemod_to_file(
     if output_tree.deep_equals(source_tree):
         return False
 
-    diff = create_diff(source_tree, output_tree)
+    diff = create_diff_from_tree(source_tree, output_tree)
     change_set = ChangeSet(
         str(file_context.file_path.relative_to(base_directory)),
         diff,

src/codemodder/context.py

@@ -6,7 +6,6 @@
 
 from codemodder.change import ChangeSet
 from codemodder.dependency import Dependency
-from codemodder.dependency_management import DependencyManager
 from codemodder.executor import CodemodExecutorWrapper
 from codemodder.file_context import FileContext
 from codemodder.logging import logger, log_list
@@ -102,6 +101,9 @@ def process_dependencies(self, codemod_id: str):
             )
             return
 
+        # pylint: disable-next=cyclic-import
+        from codemodder.dependency_management import DependencyManager
+
         dm = DependencyManager(dependencies_store, self.directory)
         if (changeset := dm.write(list(dependencies), self.dry_run)) is not None:
             self.add_results(codemod_id, [changeset])

src/codemodder/dependency_management/base_dependency_writer.py

@@ -2,9 +2,10 @@
 from pathlib import Path
 from typing import Optional
 from codemodder.project_analysis.file_parsers.package_store import PackageStore
-from codemodder.change import ChangeSet
+from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
 from codemodder.dependency import Dependency
 from packaging.requirements import Requirement
+from typing import Callable, Union, List
 
 
 class DependencyWriter(metaclass=ABCMeta):
@@ -38,3 +39,25 @@ def add(self, dependencies: list[Dependency]) -> list[Dependency]:
                 self.dependency_store.dependencies.add(requirement)
                 new.append(new_dep)
         return new
+
+    def build_changes(
+        self,
+        dependencies: list[Dependency],
+        line_number_strategy: Callable,
+        strategy_arg: Union[int, List[str], List[int]],
+    ) -> list[Change]:
+        return [
+            Change(
+                lineNumber=line_number_strategy(strategy_arg, 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)
+        ]

src/codemodder/dependency_management/dependency_manager.py

@@ -5,7 +5,14 @@
     RequirementsTxtWriter,
 )
 from codemodder.dependency_management.pyproject_writer import PyprojectWriter
-from codemodder.project_analysis.file_parsers.package_store import PackageStore
+from codemodder.dependency_management.setup_py_writer import (
+    SetupPyWriter,
+)
+
+from codemodder.project_analysis.file_parsers.package_store import (
+    PackageStore,
+    FileType,
+)
 from pathlib import Path
 
 
@@ -24,14 +31,16 @@
         Write `dependencies` to the appropriate location in the project.
         """
         match self.dependencies_store.type:
-            case "requirements.txt":
+            case FileType.REQ_TXT:
                 return RequirementsTxtWriter(
                     self.dependencies_store, self.parent_directory
                 ).write(dependencies, dry_run)
-            case "pyproject.toml":
+            case FileType.TOML:
                 return PyprojectWriter(
                     self.dependencies_store, self.parent_directory
                 ).write(dependencies, dry_run)
-            case "setup.py":
-                pass
+            case FileType.SETUP_PY:
+                return SetupPyWriter(
+                    self.dependencies_store, self.parent_directory
+                ).write(dependencies, dry_run)
         return None

src/codemodder/dependency_management/pyproject_writer.py

@@ -2,9 +2,14 @@
 from typing import Optional
 from copy import deepcopy
 from codemodder.dependency import Dependency
-from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
+from codemodder.change import ChangeSet
 from codemodder.dependency_management.base_dependency_writer import DependencyWriter
 from codemodder.diff import create_diff_and_linenums
+from codemodder.logging import logger
+
+
+def added_line_nums_strategy(lines, i):
+    return lines[i]
 
 
 class PyprojectWriter(DependencyWriter):
@@ -19,6 +24,7 @@ def add_to_file(
                 [f"{dep.requirement}" for dep in dependencies]
             )
         except tomlkit.exceptions.NonExistentKey:
+            logger.debug("Unable to add dependencies to pyproject.toml file.")
             return None
 
         diff, added_line_nums = create_diff_and_linenums(
@@ -29,21 +35,9 @@ def add_to_file(
             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)
-        ]
+        changes = self.build_changes(
+            dependencies, added_line_nums_strategy, added_line_nums
+        )
         return ChangeSet(
             str(self.path.relative_to(self.parent_directory)),
             diff,

src/codemodder/dependency_management/requirements_txt_writer.py

@@ -1,10 +1,14 @@
 from typing import Optional
 from codemodder.dependency_management.base_dependency_writer import DependencyWriter
-from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
+from codemodder.change import ChangeSet
 from codemodder.diff import create_diff
 from codemodder.dependency import Dependency
 
 
+def original_lines_strategy(original_lines, i):
+    return len(original_lines) + i + 1
+
+
 class RequirementsTxtWriter(DependencyWriter):
     def add_to_file(
         self, dependencies: list[Dependency], dry_run: bool = False
@@ -23,21 +27,9 @@ def add_to_file(
             with open(self.path, "w", encoding="utf-8") as f:
                 f.writelines(updated_lines)
 
-        changes = [
-            Change(
-                lineNumber=len(original_lines) + i + 1,
-                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)
-        ]
+        changes = self.build_changes(
+            dependencies, original_lines_strategy, original_lines
+        )
         return ChangeSet(
             str(self.path.relative_to(self.parent_directory)),
             diff,

src/codemodder/dependency_management/setup_py_codemod.py → src/codemodder/dependency_management/setup_py_writer.py

@@ -1,12 +1,59 @@
 import libcst as cst
 from libcst.codemod import CodemodContext
 from libcst import matchers
+from typing import Optional
 from codemodder.codemods.api import BaseCodemod
 from codemodder.codemods.base_codemod import ReviewGuidance
 from codemodder.codemods.utils import is_setup_py_file
 from codemodder.codemods.utils_mixin import NameResolutionMixin
 from codemodder.file_context import FileContext
 from packaging.requirements import Requirement
+from codemodder.dependency import Dependency
+from codemodder.dependency_management.base_dependency_writer import DependencyWriter
+from codemodder.change import ChangeSet
+from codemodder.diff import create_diff_from_tree
+
+
+def fixed_line_number_strategy(line_num_changed, _):
+    return line_num_changed
+
+
+class SetupPyWriter(DependencyWriter):
+    def add_to_file(
+        self, dependencies: list[Dependency], dry_run: bool = False
+    ) -> Optional[ChangeSet]:
+        input_tree = self._parse_file()
+        wrapper = cst.MetadataWrapper(input_tree)
+        file_context = FileContext(self.parent_directory, self.path, [], [], [])
+
+        codemod = SetupPyAddDependencies(
+            CodemodContext(wrapper=wrapper),
+            file_context,
+            dependencies=[dep.requirement for dep in dependencies],
+        )
+
+        output_tree = codemod.transform_module(input_tree)
+        if codemod.line_num_changed is None:
+            return None
+
+        diff = create_diff_from_tree(input_tree, output_tree)
+
+        if not dry_run:
+            with open(self.path, "w", encoding="utf-8") as f:
+                f.write(output_tree.code)
+
+        changes = self.build_changes(
+            dependencies, fixed_line_number_strategy, codemod.line_num_changed
+        )
+        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 cst.parse_module(f.read())
 
 
 class SetupPyAddDependencies(BaseCodemod, NameResolutionMixin):
@@ -26,6 +73,7 @@ def __init__(
         NameResolutionMixin.__init__(self)
         self.filename = self.file_context.file_path
         self.dependencies = dependencies
+        self.line_num_changed = None
 
     def visit_Module(self, _: cst.Module) -> bool:
         """
@@ -54,6 +102,14 @@ def replace_arg(self, original_node: cst.Call):
         return new_args
 
     def add_dependencies_to_arg(self, arg: cst.Arg) -> cst.Arg:
+        if not arg.value.elements:
+            # If there are no current dependencies, don't do anything
+            return arg
+
+        # we add the new dependencies in the same line as the last
+        # dependency listed in install_requires
+        self.line_num_changed = self.lineno_for_node(arg.value.elements[-1]) - 1
+
         new_dependencies = [
             cst.Element(value=cst.SimpleString(value=f'"{str(dep)}"'))
             for dep in self.dependencies

src/codemodder/diff.py

@@ -1,11 +1,22 @@
 import difflib
+import libcst as cst
 
 
 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_from_tree(original_tree: cst.Module, new_tree: cst.Module) -> str:
+    """
+    Create a diff between the original and output trees.
+    """
+    return create_diff(
+        original_tree.code.splitlines(keepends=True),
+        new_tree.code.splitlines(keepends=True),
+    )
+
+
 def create_diff_and_linenums(
     original_lines: list[str], new_lines: list[str]
 ) -> tuple[str, list[int]]:

src/codemodder/project_analysis/file_parsers/base_parser.py

@@ -12,7 +12,7 @@ def __init__(self, parent_directory: Path):
 
     @property
     @abstractmethod
-    def file_name(self):
+    def file_type(self):
         ...  # pragma: no cover
 
     def _parse_dependencies(self, dependencies: List[str]):
@@ -28,7 +28,7 @@ def _parse_file(self, file: Path):
         ...  # pragma: no cover
 
     def find_file_locations(self) -> List[Path]:
-        return list(Path(self.parent_directory).rglob(self.file_name))
+        return list(Path(self.parent_directory).rglob(self.file_type.value))
 
     def parse(self) -> list[PackageStore]:
         """

src/codemodder/project_analysis/file_parsers/package_store.py

@@ -1,10 +1,18 @@
 from dataclasses import dataclass
+from enum import Enum
 from packaging.requirements import Requirement
 
 
+class FileType(Enum):
+    REQ_TXT = "requirements.txt"
+    TOML = "pyproject.toml"
+    SETUP_PY = "setup.py"
+    SETUP_CFG = "setup.cfg"
+
+
 @dataclass
 class PackageStore:
-    type: str
+    type: FileType
     file: str
     dependencies: set[Requirement]
     py_versions: list[str]

src/codemodder/project_analysis/file_parsers/pyproject_toml_file_parser.py

@@ -1,4 +1,7 @@
-from codemodder.project_analysis.file_parsers.package_store import PackageStore
+from codemodder.project_analysis.file_parsers.package_store import (
+    PackageStore,
+    FileType,
+)
 from pathlib import Path
 import toml
 
@@ -7,8 +10,8 @@
 
 class PyprojectTomlParser(BaseParser):
     @property
-    def file_name(self):
-        return "pyproject.toml"
+    def file_type(self):
+        return FileType.TOML
 
     def _parse_dependencies_from_toml(self, toml_data: dict):
         # todo: handle cases for
@@ -27,7 +30,7 @@ def _parse_file(self, file: Path):
         # todo: handle no "project" in data
 
         return PackageStore(
-            type=self.file_name,
+            type=self.file_type,
             file=str(file),
             dependencies=set(self._parse_dependencies_from_toml(data)),
             py_versions=self._parse_py_versions(data),