create RequirementsWriter, call PRM

src/codemodder/codemodder.py

@@ -93,14 +93,15 @@ def apply_codemod_to_file(
     return True
 
 
+# pylint: disable-next=too-many-arguments
 def process_file(
     idx: int,
     file_path: Path,
     base_directory: Path,
     codemod,
     results: ResultSet,
     cli_args,
-):  # pylint: disable=too-many-arguments
+) -> FileContext:
     logger.debug("scanning file %s", file_path)
     if idx and idx % 100 == 0:
         logger.info("scanned %s files...", idx)  # pragma: no cover
@@ -262,8 +263,8 @@ def run(original_args) -> int:
         codemod_registry,
         repo_manager,
     )
-    # todo: enable when ready
-    # repo_manager.package_stores
+
+    repo_manager.parse_project()
 
     # TODO: this should be a method of CodemodExecutionContext
     codemods_to_run = codemod_registry.match_codemods(

src/codemodder/context.py

@@ -87,20 +87,23 @@ def get_failed_files(self):
         )
 
     def process_dependencies(self, codemod_id: str):
+        """Write the dependencies a codemod added to the appropriate dependency
+        file in the project.
+        """
         dependencies = self.dependencies.get(codemod_id)
         if not dependencies:
             return
 
-        dm = DependencyManager(self.directory)
-        if not dm.found_dependency_file:
+        dependencies_store = self.repo_manager.dependencies_store
+        if dependencies_store is None:
             logger.info(
                 "unable to write dependencies for %s: no dependency file found",
                 codemod_id,
             )
             return
 
-        dm.add(list(dependencies))
-        if (changeset := dm.write(self.dry_run)) is not None:
+        dm = DependencyManager(dependencies_store, self.directory)
+        if (changeset := dm.write(list(dependencies), self.dry_run)) is not None:
             self.add_results(codemod_id, [changeset])
 
     def add_description(self, codemod: CodemodExecutorWrapper):

src/codemodder/dependency_management/__init__.py

@@ -1 +1 @@
-from .dependency_manager import DependencyManager, Requirement
+from .dependency_manager import DependencyManager

src/codemodder/dependency_management/base_dependency_writer.py

@@ -1,19 +1,32 @@
 from abc import ABCMeta, abstractmethod
 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.dependency import Requirement
+from codemodder.dependency import Dependency
+from packaging.requirements import Requirement
 
 
 class DependencyWriter(metaclass=ABCMeta):
-    path: Path
+    dependency_store: PackageStore
 
-    def __init__(self, path: str | Path):
-        self.path = Path(path)
+    def __init__(self, dependency_store: PackageStore, parent_directory: Path):
+        self.dependency_store = dependency_store
+        self.path = Path(dependency_store.file)
+        self.parent_directory = parent_directory
 
     @abstractmethod
     def write(
         self, dependencies: list[Requirement], dry_run: bool = False
     ) -> Optional[ChangeSet]:
         pass
+
+    def add(self, dependencies: list[Dependency]) -> Optional[list[Dependency]]:
+        """add any number of dependencies to the end of list of dependencies."""
+        new = []
+        for new_dep in dependencies:
+            requirement: Requirement = new_dep.requirement
+            if requirement not in self.dependency_store.dependencies:
+                self.dependency_store.dependencies.append(requirement)
+                new.append(new_dep)
+        return new

src/codemodder/dependency_management/dependency_manager.py

@@ -1,110 +1,32 @@
-from functools import cached_property
-from pathlib import Path
 from typing import Optional
-
-from packaging.requirements import Requirement
-
-from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
-from codemodder.diff import create_diff
+from codemodder.change import ChangeSet
 from codemodder.dependency import Dependency
+from codemodder.dependency_management.requirements_txt_writer import (
+    RequirementsTxtWriter,
+)
+from codemodder.project_analysis.file_parsers.package_store import PackageStore
+from pathlib import Path
 
 
 class DependencyManager:
+    dependencies_store: PackageStore
     parent_directory: Path
-    _lines: list[str]
-    _new_requirements: list[Dependency]
 
-    def __init__(self, parent_directory: Path):
+    def __init__(self, dependencies_store: PackageStore, parent_directory: Path):
+        self.dependencies_store = dependencies_store
         self.parent_directory = parent_directory
-        self.dependency_file_changed = False
-        self._lines = []
-        self._new_requirements = []
-
-    @property
-    def new_requirements(self) -> list[str]:
-        return [str(x.requirement) for x in self._new_requirements]
 
-    def add(self, dependencies: list[Dependency]):
-        """add any number of dependencies to the end of list of dependencies."""
-        for dep in dependencies:
-            if dep.requirement.name not in self.dependencies:
-                self.dependencies.update({dep.requirement.name: dep.requirement})
-                self._new_requirements.append(dep)
-
-    def write(self, dry_run: bool = False) -> Optional[ChangeSet]:
+    def write(
+        self, dependencies: list[Dependency], dry_run: bool = False
+    ) -> Optional[ChangeSet]:
         """
-        Write the updated dependency files if any changes were made.
+        Write `dependencies` to the appropriate location in the project.
         """
-        if not (self.dependency_file and self._new_requirements):
-            return None
-
-        original_lines = self._lines.copy()
-        if not original_lines[-1].endswith("\n"):
-            original_lines[-1] += "\n"
-
-        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(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(self._new_requirements)
-        ]
-
-        if not dry_run:
-            with open(self.dependency_file, "w", encoding="utf-8") as f:
-                f.writelines(original_lines)
-                f.writelines(requirement_lines)
-
-        self.dependency_file_changed = True
-        return ChangeSet(
-            str(self.dependency_file.relative_to(self.parent_directory)),
-            diff,
-            changes=changes,
-        )
-
-    @property
-    def found_dependency_file(self) -> bool:
-        return self.dependency_file is not None
-
-    @cached_property
-    def dependency_file(self) -> Optional[Path]:
-        try:
-            # For now for simplicity only return the first file
-            return next(Path(self.parent_directory).rglob("requirements.txt"))
-        except StopIteration:
-            pass
+        match self.dependencies_store.type:
+            case "requirements.txt":
+                return RequirementsTxtWriter(
+                    self.dependencies_store, self.parent_directory
+                ).write(dependencies, dry_run)
+            case "setup.py":
+                pass
         return None
-
-    @cached_property
-    def dependencies(self) -> dict[str, Requirement]:
-        """
-        Extract list of dependencies from requirements.txt file.
-        Same order of requirements is maintained, no alphabetical sorting is done.
-        """
-        if not self.dependency_file:
-            return {}
-
-        with open(self.dependency_file, "r", encoding="utf-8") as f:
-            self._lines = f.readlines()
-
-        return {
-            requirement.name: requirement
-            for x in self._lines
-            # Skip empty lines and comments
-            if (line := x.strip())
-            and not line.startswith("#")
-            and (requirement := Requirement(line))
-        }

src/codemodder/dependency_management/requirements_txt_writer.py

@@ -1,8 +1,57 @@
+from typing import Optional
 from codemodder.dependency_management.base_dependency_writer import DependencyWriter
+from codemodder.change import Action, Change, ChangeSet, PackageAction, Result
+from packaging.requirements import Requirement
+from codemodder.diff import create_diff
 
 
 class RequirementsTxtWriter(DependencyWriter):
     def write(
         self, dependencies: list[Requirement], dry_run: bool = False
     ) -> Optional[ChangeSet]:
-        pass
+        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):
+        original_lines = self._parse_file()
+        updated_lines = original_lines.copy()
+        if not original_lines[-1].endswith("\n"):
+            updated_lines[-1] += "\n"
+
+        requirement_lines = [f"{dep.requirement}\n" for dep in dependencies]
+        updated_lines += requirement_lines
+
+        diff = create_diff(original_lines, 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)
+        ]
+
+        if not dry_run:
+            with open(self.path, "w", encoding="utf-8") as f:
+                f.writelines(original_lines)
+                f.writelines(requirement_lines)
+
+        return ChangeSet(
+            str(self.path.relative_to(self.parent_directory)),
+            diff,
+            changes=changes,
+        )
+
+    def _parse_file(self):
+        with open(self.path, "r", encoding="utf-8") as f:
+            return f.readlines()

src/codemodder/project_analysis/python_repo_manager.py

@@ -1,5 +1,6 @@
 from functools import cached_property
 from pathlib import Path
+from typing import Optional
 from codemodder.project_analysis.file_parsers import (
     RequirementsTxtParser,
     PyprojectTomlParser,
@@ -19,10 +20,23 @@ def __init__(self, parent_directory: Path):
             SetupPyParser,
         ]
 
+    @cached_property
+    def dependencies_store(self) -> Optional[PackageStore]:
+        """The location where to write new dependencies for project.
+        For now just pick the first store found with order given by _potential_stores.
+        """
+        if self.package_stores:
+            return self.package_stores[0]
+        return None
+
     @cached_property
     def package_stores(self) -> list[PackageStore]:
         return self._parse_all_stores()
 
+    def parse_project(self) -> list[PackageStore]:
+        """Wrapper around cached-property for clarity when calling it the first time."""
+        return self.package_stores
+
     def _parse_all_stores(self) -> list[PackageStore]:
         discovered_pkg_stores: list[PackageStore] = []
         for store in self._potential_stores: