Support parsing/writing pyproject.toml with poetry (#646)

* pyproject parser can parse poetry

* pyproject writer can add to poetry dependencies

* test pyproject writer with both [project] and poetry

* handle pyproject writer w/o poetry dependencies declared

* requirement parser can handle ^ caret syntax

* add newline

Author: clavedeluna

src/codemodder/dependency_management/pyproject_writer.py

@@ -21,13 +21,34 @@ def add_to_file(
         pyproject = self._parse_file()
         original = deepcopy(pyproject)
 
-        try:
-            pyproject["project"]["dependencies"].extend(
-                [f"{dep.requirement}" for dep in dependencies]
-            )
-        except tomlkit.exceptions.NonExistentKey:
-            logger.debug("Unable to add dependencies to pyproject.toml file.")
-            return None
+        if poetry_data := pyproject.get("tool", {}).get("poetry", {}):
+            add_newline = False
+            # 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())
+
+        else:
+            try:
+                pyproject["project"]["dependencies"].extend(
+                    [f"{dep.requirement}" for dep in dependencies]
+                )
+            except tomlkit.exceptions.NonExistentKey:
+                logger.debug("Unable to add dependencies to pyproject.toml file.")
+                return None
 
         diff, added_line_nums = create_diff_and_linenums(
             tomlkit.dumps(original).split("\n"), tomlkit.dumps(pyproject).split("\n")

src/codemodder/project_analysis/file_parsers/package_store.py

@@ -2,6 +2,8 @@
 from enum import Enum
 from pathlib import Path
 
+from packaging.requirements import InvalidRequirement
+
 from codemodder.dependency import Requirement
 
 
@@ -29,10 +31,38 @@ def __init__(
         self.type = type
         self.file = file
         self.dependencies = {
-            dep if isinstance(dep, Requirement) else Requirement(dep)
-            for dep in dependencies
+            x for x in {parse_requirement(dep) for dep in dependencies} if x
         }
         self.py_versions = py_versions
 
     def has_requirement(self, requirement: Requirement) -> bool:
         return requirement.name in {dep.name for dep in self.dependencies}
+
+
+def parse_requirement(requirement: str | Requirement) -> Requirement:
+    match requirement:
+        case Requirement():
+            return requirement
+        case _:
+            try:
+                return Requirement(convert_py_version(requirement))
+            except (InvalidRequirement, ValueError):
+                return None
+
+
+def convert_py_version(py_version: str) -> str:
+    """
+    Convert any dependency with ^ syntax to equivalent >=, < syntax.
+    packaging.requirements does not support parsing dependencies with `^` syntax
+    which is common in Poetry. For our internal representation we will convert it to
+
+    `pandas^1.2.31` is equivalent to `pandas>=1.2.3,< 2.0.0`
+    """
+    if "^" in py_version:
+        try:
+            name, version = py_version.split("^")
+            next_major_version = int(version.strip()[0]) + 1
+            return f"{name}>={version},<{next_major_version}.0.0"
+        except Exception:
+            raise ValueError
+    return py_version

src/codemodder/project_analysis/file_parsers/pyproject_toml_file_parser.py

@@ -16,17 +16,35 @@ def file_type(self):
         return FileType.TOML
 
     def _parse_file(self, file: Path) -> PackageStore | None:
+        """
+        Parse a pyproject.toml file which may or may not use poetry.
+        """
         data = toml.load(file)
+        project_data = data.get("project")
+        poetry_data = data.get("tool", {}).get("poetry")
 
-        if not (project := data.get("project")):
+        if not project_data and not poetry_data:
             return None
 
-        dependencies = project.get("dependencies", [])
-        version = project.get("requires-python", None)
+        version = None
+        project_dependencies, poetry_dependencies = [], []
+        if project_data:
+            project_dependencies = project_data.get("dependencies", [])
+            version = project_data.get("requires-python", None)
+
+        if poetry_data:
+            poetry_dependencies = [
+                f"{name}{version}"
+                for name, version in poetry_data.get("dependencies", {}).items()
+                if name != "python"
+            ]
+
+            # In poetry python version is declared within `[tool.poetry.dependencies]`
+            version = poetry_data.get("dependencies", {}).get("python", None)
 
         return PackageStore(
             type=self.file_type,
             file=file,
-            dependencies=set(dependencies),
+            dependencies=set(project_dependencies + poetry_dependencies),
             py_versions=[version] if version else [],
         )

tests/dependency_management/test_pyproject_writer.py

@@ -36,7 +36,7 @@ def test_update_pyproject_dependencies(tmpdir, dry_run):
     pyproject_toml.write(dedent(orig_pyproject))
 
     store = PackageStore(
-        type=FileType.REQ_TXT,
+        type=FileType.TOML,
         file=pyproject_toml,
         dependencies=set(),
         py_versions=[">=3.10.0"],
@@ -127,7 +127,7 @@ def test_add_same_dependency_only_once(tmpdir):
     pyproject_toml.write(dedent(orig_pyproject))
 
     store = PackageStore(
-        type=FileType.REQ_TXT,
+        type=FileType.TOML,
         file=pyproject_toml,
         dependencies=set(),
         py_versions=[">=3.10.0"],
@@ -181,7 +181,7 @@ def test_dont_add_existing_dependency(tmpdir):
     pyproject_toml.write(dedent(orig_pyproject))
 
     store = PackageStore(
-        type=FileType.REQ_TXT,
+        type=FileType.TOML,
         file=pyproject_toml,
         dependencies=set([Security.requirement]),
         py_versions=[">=3.10.0"],
@@ -207,7 +207,7 @@ def test_pyproject_no_dependencies(tmpdir):
     pyproject_toml.write(dedent(orig_pyproject))
 
     store = PackageStore(
-        type=FileType.REQ_TXT,
+        type=FileType.TOML,
         file=pyproject_toml,
         dependencies=set(),
         py_versions=[">=3.10.0"],
@@ -219,3 +219,192 @@ def test_pyproject_no_dependencies(tmpdir):
     writer.write(dependencies)
 
     assert pyproject_toml.read() == dedent(orig_pyproject)
+
+
+def test_pyproject_poetry(tmpdir):
+    orig_pyproject = """\
+        [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"
+    """
+
+    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)
+
+    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"
+        {Security.requirement.name} = "{ str(Security.requirement.specifier)}"
+        {DefusedXML.requirement.name} = "{ str(DefusedXML.requirement.specifier)}"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)
+
+
+def test_pyproject_poetry_existing_dependency(tmpdir):
+    orig_pyproject = """\
+        [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"
+        defusedxml = "^0.7"
+    """
+
+    pyproject_toml = tmpdir.join("pyproject.toml")
+    pyproject_toml.write(dedent(orig_pyproject))
+
+    store = PackageStore(
+        type=FileType.TOML,
+        file=pyproject_toml,
+        dependencies=set([DefusedXML.requirement]),
+        py_versions=["~=3.11.0"],
+    )
+
+    writer = PyprojectWriter(store, tmpdir)
+    dependencies = [DefusedXML]
+    writer.write(dependencies)
+
+    assert pyproject_toml.read() == dedent(orig_pyproject)
+
+
+def test_pyproject_poetry_no_deps_section(tmpdir):
+    orig_pyproject = """\
+        [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"
+    """
+
+    pyproject_toml = tmpdir.join("pyproject.toml")
+    pyproject_toml.write(dedent(orig_pyproject))
+
+    store = PackageStore(
+        type=FileType.TOML,
+        file=pyproject_toml,
+        dependencies=set(),
+        py_versions=[],
+    )
+
+    writer = PyprojectWriter(store, tmpdir)
+    dependencies = [Security, DefusedXML]
+    writer.write(dependencies)
+
+    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"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)
+
+
+def test_pyproject_poetry_no_declared_deps(tmpdir):
+    orig_pyproject = """\
+        [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"
+    """
+
+    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)
+
+    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"
+        {Security.requirement.name} = "{str(Security.requirement.specifier)}"
+        {DefusedXML.requirement.name} = "{str(DefusedXML.requirement.specifier)}"
+    """
+
+    assert pyproject_toml.read() == dedent(updated_pyproject)