codemod plain cst

src/codemodder/codemods/utils.py

@@ -113,6 +113,11 @@ 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 get_call_name(call: cst.Call) -> str:
     """
     Extracts the full name from a function call

src/codemodder/dependency_management/setup_py_codemod.py

@@ -0,0 +1,37 @@
+import libcst as cst
+from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
+from codemodder.codemods.utils import is_setup_py_file
+from codemodder.codemods.utils_mixin import NameResolutionMixin
+
+
+class SetupPyAddDependencies(VisitorBasedCodemodCommand, NameResolutionMixin):
+    def __init__(self, context: CodemodContext, dependencies):
+        """
+        :param dependencies:
+        """
+        super().__init__(context)
+        self.filename = self.context.filename
+        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
+
+        # todo: add self.dependencies to install_requires arg
+        breakpoint()
+        return updated_node
+
+
+# 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)

tests/dependency_management/test_setup_py_codemod.py

@@ -0,0 +1,41 @@
+from codemodder.dependency_management.setup_py_codemod import SetupPyAddDependencies
+from libcst.codemod import CodemodTest, CodemodContext
+from packaging.requirements import Requirement
+
+TEST_DEPENDENCIES = [Requirement("defusedxml==0.7.1"), Requirement("security~=1.2.0")]
+
+
+class TestSetupPyCodemod(CodemodTest):
+    TRANSFORM = SetupPyAddDependencies
+    CONTEXT = CodemodContext(filename="pkg/setup.py")
+
+    def test_setup_call(self):
+        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 = ""
+
+        self.assertCodemod(
+            before, after, TEST_DEPENDENCIES, context_override=self.CONTEXT
+        )
+
+    # def test_different_setup_call(self):
+    # test does not call install_requires
+    # test with no dependencies inside install_requires

tests/samples/pkg_w_setuppy/setup.py

@@ -0,0 +1,25 @@
+from os import path
+from setuptools import find_packages, setup
+
+root_dir = path.abspath(path.dirname(__file__))
+
+print(root_dir)
+
+setup(
+    name="test pkg",
+    description="testing",
+    long_description="...",
+    # The project's main homepage.
+    # Author details
+    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={},
+)

tests/samples/pkg_w_setuppy/src/sample/hello.py

@@ -0,0 +1 @@
+print("hello world")