use base codemod

src/codemodder/codemods/utils.py

@@ -113,9 +113,8 @@ def is_django_settings_file(file_path: Path):
     return False
 
 
-def is_setup_py_file(file_path: str):
-    name = Path(file_path).name
-    return name == "setup.py"
+def is_setup_py_file(file_path: Path):
+    return file_path.name == "setup.py"
 
 
 def get_call_name(call: cst.Call) -> str:

src/codemodder/dependency_management/setup_py_codemod.py

@@ -1,16 +1,26 @@
 import libcst as cst
-from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
+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
 
 
-class SetupPyAddDependencies(VisitorBasedCodemodCommand, NameResolutionMixin):
-    def __init__(self, context: CodemodContext, dependencies):
-        """
-        :param dependencies:
-        """
-        super().__init__(context)
-        self.filename = self.context.filename
+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
+    ):
+        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:
@@ -24,14 +34,35 @@ def leave_Call(self, original_node: cst.Call, updated_node: cst.Call):
         if true_name != "setuptools.setup":
             return original_node
 
-        # todo: add self.dependencies to install_requires arg
-        breakpoint()
-        return updated_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
 
-# filename = "tests/samples/pkg_w_setuppy/setup.py"
-# with open(filename, "r", encoding="utf-8") as f:
-#     source_tree = cst.parse_module(f.read())
-#
-# codemod = SetupPyAddDependencies(CodemodContext(filename=filename), ["dep1"])
-# codemod.transform_module(source_tree)
+    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/dependency_management/test_setup_py_codemod.py

@@ -1,15 +1,28 @@
+import pytest
+import libcst as cst
 from codemodder.dependency_management.setup_py_codemod import SetupPyAddDependencies
-from libcst.codemod import CodemodTest, CodemodContext
+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(CodemodTest):
-    TRANSFORM = SetupPyAddDependencies
-    CONTEXT = CodemodContext(filename="pkg/setup.py")
+class TestSetupPyCodemod(BaseCodemodTest):
+    codemod = SetupPyAddDependencies
 
-    def test_setup_call(self):
+    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(
@@ -24,18 +37,182 @@ def test_setup_call(self):
                 "protobuf>=3.12,<3.18; python_version < '3'",
                 "protobuf>=3.12,<4; python_version >= '3'",
                 "psutil>=5.7,<6",
-                "requests>=2.4.2,<3",
+                "requests>=2.4.2,<3"
             ],
             entry_points={},
         )
         """
 
-        after = ""
+        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)
 
-        self.assertCodemod(
-            before, after, TEST_DEPENDENCIES, context_override=self.CONTEXT
+    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_different_setup_call(self):
-    # test does not call install_requires
-    # test with no dependencies inside install_requires
+    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)