Add requirements to setup.py codemod

src/codemodder/codemods/utils.py

@@ -113,6 +113,10 @@ def is_django_settings_file(file_path: Path):
     return False
 
 
+def is_setup_py_file(file_path: Path):
+    return file_path.name == "setup.py"
+
+
 def get_call_name(call: cst.Call) -> str:
     """
     Extracts the full name from a function call

src/codemodder/context.py

@@ -6,7 +6,7 @@
 
 from codemodder.change import ChangeSet
 from codemodder.dependency import Dependency
-from codemodder.dependency_manager import DependencyManager
+from codemodder.dependency_management import DependencyManager
 from codemodder.executor import CodemodExecutorWrapper
 from codemodder.file_context import FileContext
 from codemodder.logging import logger, log_list

src/codemodder/dependency_management/__init__.py

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

src/codemodder/dependency_management/setup_py_codemod.py

@@ -0,0 +1,72 @@
+import libcst as cst
+from libcst.codemod import CodemodContext
+from libcst import matchers
+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
+
+
+class SetupPyAddDependencies(BaseCodemod, NameResolutionMixin):
+    NAME = "setup-py-add-dependencies"
+    REVIEW_GUIDANCE = ReviewGuidance.MERGE_WITHOUT_REVIEW
+    SUMMARY = "Add Dependencies to `setup.py` `install_requires`"
+    DESCRIPTION = SUMMARY
+    REFERENCES: list = []
+
+    def __init__(
+        self,
+        codemod_context: CodemodContext,
+        file_context: FileContext,
+        dependencies: list[Requirement],
+    ):
+        BaseCodemod.__init__(self, codemod_context, file_context)
+        NameResolutionMixin.__init__(self)
+        self.filename = self.file_context.file_path
+        self.dependencies = dependencies
+
+    def visit_Module(self, _: cst.Module) -> bool:
+        """
+        Only visit module with this codemod if it's a setup.py file.
+        """
+        return is_setup_py_file(self.filename)
+
+    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
+        true_name = self.find_base_name(original_node.func)
+        if true_name != "setuptools.setup":
+            return original_node
+
+        new_args = self.replace_arg(original_node)
+        return self.update_arg_target(updated_node, new_args)
+
+    def replace_arg(self, original_node: cst.Call):
+        new_args = []
+        for arg in original_node.args:
+            if matchers.matches(
+                arg.keyword, matchers.Name("install_requires")
+            ) and matchers.matches(arg.value, matchers.List()):
+                new = self.add_dependencies_to_arg(arg)
+            else:
+                new = arg
+            new_args.append(new)
+        return new_args
+
+    def add_dependencies_to_arg(self, arg: cst.Arg) -> cst.Arg:
+        new_dependencies = [
+            cst.Element(value=cst.SimpleString(value=f'"{str(dep)}"'))
+            for dep in self.dependencies
+        ]
+        # TODO: detect if elements are separated by newline in source code.
+        return cst.Arg(
+            keyword=arg.keyword,
+            value=arg.value.with_changes(
+                elements=arg.value.elements + tuple(new_dependencies)
+            ),
+            equal=arg.equal,
+            comma=arg.comma,
+            star=arg.star,
+            whitespace_after_star=arg.whitespace_after_star,
+            whitespace_after_arg=arg.whitespace_after_arg,
+        )

tests/codemods/base_codemod_test.py

@@ -21,6 +21,14 @@ class BaseCodemodTest:
     def setup_method(self):
         self.file_context = None
 
+    def initialize_codemod(self, input_tree):
+        wrapper = cst.MetadataWrapper(input_tree)
+        codemod_instance = self.codemod(
+            CodemodContext(wrapper=wrapper),
+            self.file_context,
+        )
+        return codemod_instance
+
     def run_and_assert(self, tmpdir, input_code, expected):
         tmp_file_path = Path(tmpdir / "code.py")
         self.run_and_assert_filepath(tmpdir, tmp_file_path, input_code, expected)
@@ -41,12 +49,8 @@ def run_and_assert_filepath(self, root, file_path, input_code, expected):
             [],
             [],
         )
-        wrapper = cst.MetadataWrapper(input_tree)
-        command_instance = self.codemod(
-            CodemodContext(wrapper=wrapper),
-            self.file_context,
-        )
-        output_tree = command_instance.transform_module(input_tree)
+        codemod_instance = self.initialize_codemod(input_tree)
+        output_tree = codemod_instance.transform_module(input_tree)
 
         assert output_tree.code == dedent(expected)
 
@@ -92,12 +96,8 @@ def run_and_assert_filepath(self, root, file_path, input_code, expected):
             [],
             results,
         )
-        wrapper = cst.MetadataWrapper(input_tree)
-        command_instance = self.codemod(
-            CodemodContext(wrapper=wrapper),
-            self.file_context,
-        )
-        output_tree = command_instance.transform_module(input_tree)
+        codemod_instance = self.initialize_codemod(input_tree)
+        output_tree = codemod_instance.transform_module(input_tree)
 
         assert output_tree.code == dedent(expected)
 

tests/conftest.py

@@ -44,7 +44,9 @@ def disable_write_dependencies():
     """
     Unit tests should not write any dependency files
     """
-    dm_write = mock.patch("codemodder.dependency_manager.DependencyManager.write")
+    dm_write = mock.patch(
+        "codemodder.dependency_management.dependency_manager.DependencyManager.write"
+    )
 
     dm_write.start()
     yield

tests/test_dependency_manager.py → tests/dependency_management/test_dependency_manager.py

@@ -3,7 +3,7 @@
 import pytest
 
 from codemodder.dependency import DefusedXML, Security
-from codemodder.dependency_manager import DependencyManager, Requirement
+from codemodder.dependency_management import DependencyManager, Requirement
 
 
 @pytest.fixture(autouse=True, scope="module")

tests/dependency_management/test_setup_py_codemod.py

@@ -0,0 +1,218 @@
+import pytest
+import libcst as cst
+from codemodder.dependency_management.setup_py_codemod import SetupPyAddDependencies
+from libcst.codemod import CodemodContext
+from tests.codemods.base_codemod_test import BaseCodemodTest
+from packaging.requirements import Requirement
+from pathlib import Path
+
+TEST_DEPENDENCIES = [Requirement("defusedxml==0.7.1"), Requirement("security~=1.2.0")]
+
+
+class TestSetupPyCodemod(BaseCodemodTest):
+    codemod = SetupPyAddDependencies
+
+    def initialize_codemod(self, input_tree):
+        """This codemod is initialized with different args than other codemods."""
+        wrapper = cst.MetadataWrapper(input_tree)
+        codemod_instance = self.codemod(
+            CodemodContext(wrapper=wrapper),
+            self.file_context,
+            dependencies=TEST_DEPENDENCIES,
+        )
+        return codemod_instance
+
+    def test_setup_call(self, tmpdir):
+        before = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=[
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3"
+            ],
+            entry_points={},
+        )
+        """
+
+        after = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=[
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3", "defusedxml==0.7.1", "security~=1.2.0"
+            ],
+            entry_points={},
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, after)
+
+    def test_other_setup_func(self, tmpdir):
+        before = """
+        from something import setup
+        setup(
+            name="test pkg",
+            install_requires=[
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3"
+            ],
+            entry_points={},
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, before)
+
+    def test_not_setup_file(self, tmpdir):
+        before = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=[
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3"
+            ],
+            entry_points={},
+        )
+        """
+        tmp_file_path = Path(tmpdir / "not-setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, before)
+
+    def test_setup_call_no_install_requires(self, tmpdir):
+        before = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, before)
+
+    def test_setup_no_existing_requirements(self, tmpdir):
+        before = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=[],
+            entry_points={},
+        )
+        """
+
+        after = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=["defusedxml==0.7.1", "security~=1.2.0"],
+            entry_points={},
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, after)
+
+    def test_setup_call_bad_install_requires(self, tmpdir):
+        before = """
+        from setuptools import setup
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires="some-package",
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, before)
+
+    @pytest.mark.skip("Need to add support.")
+    def test_setup_call_requirements_separate(self, tmpdir):
+        before = """
+        from setuptools import setup
+        requirements = [
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3"
+        ]
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=requirements,
+            entry_points={},
+        )
+        """
+
+        after = """
+        from setuptools import setup
+        requirements = [
+                "protobuf>=3.12,<3.18; python_version < '3'",
+                "protobuf>=3.12,<4; python_version >= '3'",
+                "psutil>=5.7,<6",
+                "requests>=2.4.2,<3", "defusedxml==0.7.1", "security~=1.2.0"
+        ]
+        setup(
+            name="test pkg",
+            description="testing",
+            long_description="...",
+            author="Pixee",
+            packages=find_packages("src"),
+            package_dir={"": "src"},
+            python_requires=">3.6",
+            install_requires=requirements,
+            entry_points={},
+        )
+        """
+        tmp_file_path = Path(tmpdir / "setup.py")
+        self.run_and_assert_filepath(tmpdir, tmp_file_path, before, after)