dependency type stubs (#651)

* add 2 type stubs

* add typing to poetry

* fix test

* add another unit test

src/codemodder/dependency.py

@@ -17,6 +17,8 @@ class Dependency:
     oss_link: str
     package_link: str
     hashes: list[str] = field(default_factory=list)
+    # Forward reference
+    type_stubs: list["Dependency"] = field(default_factory=list)
 
     @property
     def name(self) -> str:
@@ -56,6 +58,24 @@ def __hash__(self):
     ),
     oss_link="https://github.com/wtforms/flask-wtf/",
     package_link="https://pypi.org/project/Flask-WTF/",
+    type_stubs=[
+        Dependency(
+            Requirement("types-WTForms==3.1.0.20240425"),
+            hashes=[
+                "449b6e3756b2bc70657e98d989bdbf572a25466428774be96facf9debcbf6c4e",
+                "49ffc1fe5576ea0735b763fff77e7060dd39ecc661276cbd0b47099921b3a6f2",
+            ],
+            description="""\
+                    This is a type stub package for the WTForms package.
+        """,
+            _license=License(
+                "Apache-2.0",
+                "https://opensource.org/license/apache-2-0",
+            ),
+            oss_link="https://github.com/python/typeshed",
+            package_link="https://pypi.org/project/types-WTForms/",
+        ),
+    ],
 )
 
 DefusedXML = Dependency(
@@ -74,6 +94,24 @@ def __hash__(self):
     ),
     oss_link="https://github.com/tiran/defusedxml",
     package_link="https://pypi.org/project/defusedxml/",
+    type_stubs=[
+        Dependency(
+            Requirement("types-defusedxml==0.7.0.20240218"),
+            hashes=[
+                "2b7f3c5ca14fdbe728fab0b846f5f7eb98c4bd4fd2b83d25f79e923caa790ced",
+                "05688a7724dc66ea74c4af5ca0efc554a150c329cb28c13a64902cab878d06ed",
+            ],
+            description="""\
+                This is a type stub package for the defusedxml package.
+    """,
+            _license=License(
+                "Apache-2.0",
+                "https://opensource.org/license/apache-2-0",
+            ),
+            oss_link="https://github.com/python/typeshed",
+            package_link="https://pypi.org/project/types-defusedxml/",
+        ),
+    ],
 )
 
 Security = Dependency(

src/codemodder/dependency_management/codemod_dependencies.txt

@@ -5,6 +5,8 @@
 # dependency in dependency.py. Be sure to update the version AND the hashes.
 # Run `get-hashes pkg==version` to get the hashes.
 defusedxml==0.7.1
+types-defusedxml==0.7.0.20240218
 flask-wtf==1.2.0
+types-WTForms==3.1.0.20240425
 security==1.2.1
 fickling==0.1.3

src/codemodder/dependency_management/pyproject_writer.py

@@ -9,6 +9,8 @@
 from codemodder.diff import create_diff_and_linenums
 from codemodder.logging import logger
 
+TYPE_CHECKER_LIBRARIES = ["mypy", "pyright"]
+
 
 def added_line_nums_strategy(lines, i):
     return lines[i]
@@ -21,26 +23,11 @@ def add_to_file(
         pyproject = self._parse_file()
         original = deepcopy(pyproject)
 
-        if poetry_data := pyproject.get("tool", {}).get("poetry", {}):
-            add_newline = False
+        if pyproject.get("tool", {}).get("poetry", {}):
             # It's unlikely and bad practice to declare dependencies under [project].dependencies
             # and [tool.poetry.dependencies] but if it happens, we will give priority to poetry
             # and add dependencies under its system.
-            if poetry_data.get("dependencies") is None:
-                pyproject["tool"]["poetry"].append("dependencies", {})
-                add_newline = True
-
-            for dep in dependencies:
-                try:
-                    pyproject["tool"]["poetry"]["dependencies"].append(
-                        dep.requirement.name, str(dep.requirement.specifier)
-                    )
-                except tomlkit.exceptions.KeyAlreadyPresent:
-                    pass
-
-            if add_newline:
-                pyproject["tool"]["poetry"]["dependencies"].add(tomlkit.nl())
-
+            self._update_poetry(pyproject, dependencies)
         else:
             try:
                 pyproject["project"]["dependencies"].extend(
@@ -70,3 +57,67 @@ def add_to_file(
     def _parse_file(self):
         with open(self.path, encoding="utf-8") as f:
             return tomlkit.load(f)
+
+    def _update_poetry(
+        self,
+        pyproject: tomlkit.toml_document.TOMLDocument,
+        dependencies: list[Dependency],
+    ):
+        add_newline = False
+
+        if pyproject.get("tool", {}).get("poetry", {}).get("dependencies") is None:
+            pyproject["tool"]["poetry"].update({"dependencies": {}})
+            add_newline = True
+
+        typing_location = find_typing_location(pyproject)
+
+        for dep in dependencies:
+            try:
+                pyproject["tool"]["poetry"]["dependencies"].append(
+                    dep.requirement.name, str(dep.requirement.specifier)
+                )
+            except tomlkit.exceptions.KeyAlreadyPresent:
+                pass
+
+            for type_stub_dependency in dep.type_stubs:
+                if typing_location:
+                    try:
+                        keys = typing_location.split(".")
+                        section = pyproject["tool"]["poetry"]
+                        for key in keys:
+                            section = section[key]
+                        section.append(
+                            type_stub_dependency.requirement.name,
+                            str(type_stub_dependency.requirement.specifier),
+                        )
+                    except tomlkit.exceptions.KeyAlreadyPresent:
+                        pass
+
+        if add_newline:
+            pyproject["tool"]["poetry"]["dependencies"].add(tomlkit.nl())
+
+
+def find_typing_location(pyproject):
+    """
+    Look for a typing tool declared as a dependency in project.toml
+    """
+    locations = [
+        "dependencies",
+        "test.dependencies",
+        "dev-dependencies",
+        "dev.dependencies",
+        "group.test.dependencies",
+    ]
+    poetry_section = pyproject.get("tool", {}).get("poetry", {})
+
+    for location in locations:
+        keys = location.split(".")
+        section = poetry_section
+        try:
+            for key in keys:
+                section = section[key]
+            if any(checker in section for checker in TYPE_CHECKER_LIBRARIES):
+                return location
+        except KeyError:
+            continue
+    return None

tests/dependency_management/test_pyproject_writer.py

@@ -4,7 +4,10 @@
 
 from codemodder.codetf import DiffSide
 from codemodder.dependency import DefusedXML, Security
-from codemodder.dependency_management.pyproject_writer import PyprojectWriter
+from codemodder.dependency_management.pyproject_writer import (
+    TYPE_CHECKER_LIBRARIES,
+    PyprojectWriter,
+)
 from codemodder.project_analysis.file_parsers.package_store import (
     FileType,
     PackageStore,
@@ -408,3 +411,195 @@ def test_pyproject_poetry_no_declared_deps(tmpdir):
     """
 
     assert pyproject_toml.read() == dedent(updated_pyproject)
+
+
+@pytest.mark.parametrize("type_checker", TYPE_CHECKER_LIBRARIES)
+def test_pyproject_poetry_with_type_checker_tool(tmpdir, type_checker):
+    orig_pyproject = f"""\
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+
+        [tool.poetry.dependencies]
+        python = "~=3.11.0"
+        requests = ">=2.25.1,<3.0.0"
+        pandas = "^1.2.3"
+        libcst = ">1.0"
+        {type_checker} = "==1.0"
+    """
+
+    pyproject_toml = tmpdir.join("pyproject.toml")
+    pyproject_toml.write(dedent(orig_pyproject))
+
+    store = PackageStore(
+        type=FileType.TOML,
+        file=pyproject_toml,
+        dependencies=set(),
+        py_versions=["~=3.11.0"],
+    )
+
+    writer = PyprojectWriter(store, tmpdir)
+    dependencies = [Security, DefusedXML]
+    writer.write(dependencies)
+
+    defusedxml_type_stub = DefusedXML.type_stubs[0]
+    updated_pyproject = f"""\
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+
+        [tool.poetry.dependencies]
+        python = "~=3.11.0"
+        requests = ">=2.25.1,<3.0.0"
+        pandas = "^1.2.3"
+        libcst = ">1.0"
+        {type_checker} = "==1.0"
+        {Security.requirement.name} = "{str(Security.requirement.specifier)}"
+        {DefusedXML.requirement.name} = "{str(DefusedXML.requirement.specifier)}"
+        {defusedxml_type_stub.requirement.name} = "{str(defusedxml_type_stub.requirement.specifier)}"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)
+
+
+@pytest.mark.parametrize(
+    "dependency_section",
+    [
+        "[tool.poetry.test.dependencies]",
+        "[tool.poetry.dev-dependencies]",
+        "[tool.poetry.dev.dependencies]",
+        "[tool.poetry.group.test.dependencies]",
+    ],
+)
+@pytest.mark.parametrize("type_checker", TYPE_CHECKER_LIBRARIES)
+def test_pyproject_poetry_with_type_checker_tool_without_poetry_deps_section(
+    tmpdir, type_checker, dependency_section
+):
+    orig_pyproject = f"""\
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+        
+        {dependency_section}
+        {type_checker} = "==1.0"
+    """
+
+    pyproject_toml = tmpdir.join("pyproject.toml")
+    pyproject_toml.write(dedent(orig_pyproject))
+
+    store = PackageStore(
+        type=FileType.TOML,
+        file=pyproject_toml,
+        dependencies=set(),
+        py_versions=["~=3.11.0"],
+    )
+
+    writer = PyprojectWriter(store, tmpdir)
+    dependencies = [Security, DefusedXML]
+    writer.write(dependencies)
+
+    defusedxml_type_stub = DefusedXML.type_stubs[0]
+    updated_pyproject = f"""\
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [tool.poetry.dependencies]
+        {Security.requirement.name} = "{str(Security.requirement.specifier)}"
+        {DefusedXML.requirement.name} = "{str(DefusedXML.requirement.specifier)}"
+
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+        
+        {dependency_section}
+        {type_checker} = "==1.0"
+        {defusedxml_type_stub.requirement.name} = "{str(defusedxml_type_stub.requirement.specifier)}"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)
+
+
+@pytest.mark.parametrize("type_checker", TYPE_CHECKER_LIBRARIES)
+def test_pyproject_poetry_with_type_checker_tool_multiple(tmpdir, type_checker):
+    orig_pyproject = f"""\
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+        
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [tool.poetry.dependencies]
+        python = "~=3.11.0"
+        requests = ">=2.25.1,<3.0.0"
+        pandas = "^1.2.3"
+        libcst = ">1.0"
+        
+        [tool.poetry.group.test.dependencies]
+        {type_checker} = "==1.0"
+    """
+
+    pyproject_toml = tmpdir.join("pyproject.toml")
+    pyproject_toml.write(dedent(orig_pyproject))
+
+    store = PackageStore(
+        type=FileType.TOML,
+        file=pyproject_toml,
+        dependencies=set(),
+        py_versions=["~=3.11.0"],
+    )
+
+    writer = PyprojectWriter(store, tmpdir)
+    dependencies = [Security, DefusedXML]
+    writer.write(dependencies)
+
+    defusedxml_type_stub = DefusedXML.type_stubs[0]
+    updated_pyproject = f"""\
+        [build-system]
+        requires = ["poetry-core>=1.0.0"]
+        build-backend = "poetry.core.masonry.api"
+        
+        [tool.poetry]
+        name = "example-project"
+        version = "0.1.0"
+        description = "An example project to demonstrate Poetry configuration."
+        authors = ["Your Name <[email protected]>"]
+
+        [tool.poetry.dependencies]
+        python = "~=3.11.0"
+        requests = ">=2.25.1,<3.0.0"
+        pandas = "^1.2.3"
+        libcst = ">1.0"
+        {Security.requirement.name} = "{str(Security.requirement.specifier)}"
+        {DefusedXML.requirement.name} = "{str(DefusedXML.requirement.specifier)}"
+
+        [tool.poetry.group.test.dependencies]
+        {type_checker} = "==1.0"
+        {defusedxml_type_stub.requirement.name} = "{str(defusedxml_type_stub.requirement.specifier)}"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)