feat(cli): support PEP 735 on self remove (#10130)

Author: finswimmer

src/poetry/console/commands/self/self_command.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
+import typing
+
 from pathlib import Path
 from typing import TYPE_CHECKING
+from typing import Any
 
 from packaging.utils import canonicalize_name
 from poetry.core.packages.dependency import Dependency
@@ -62,25 +65,35 @@ def activated_groups(self) -> set[NormalizedName]:
 
     def generate_system_pyproject(self) -> None:
         preserved = {}
+        preserved_groups: dict[str, Any] = {}
 
         if self.system_pyproject.exists():
-            content = PyProjectTOML(self.system_pyproject).poetry_config
+            toml_file = PyProjectTOML(self.system_pyproject)
+            content = toml_file.data
 
             for key in {"group", "source"}:
-                if key in content:
-                    preserved[key] = content[key]
+                if key in toml_file.poetry_config:
+                    preserved[key] = toml_file.poetry_config[key]
+
+            if "dependency-groups" in content:
+                preserved_groups = typing.cast(
+                    "dict[str, Any]", content["dependency-groups"]
+                )
 
         package = ProjectPackage(name="poetry-instance", version=__version__)
         package.add_dependency(Dependency(name="poetry", constraint=f"{__version__}"))
 
         package.python_versions = ".".join(str(v) for v in self.env.version_info[:3])
 
-        content = Factory.create_pyproject_from_package(package=package)
+        content = Factory.create_legacy_pyproject_from_package(package=package)
         content["tool"]["poetry"]["package-mode"] = False  # type: ignore[index]
 
         for key in preserved:
             content["tool"]["poetry"][key] = preserved[key]  # type: ignore[index]
 
+        if preserved_groups:
+            content["dependency-groups"] = preserved_groups
+
         pyproject = PyProjectTOML(self.system_pyproject)
         pyproject.file.write(content)
 

src/poetry/factory.py

@@ -220,7 +220,7 @@ def create_package_source(
         )
 
     @classmethod
-    def create_pyproject_from_package(cls, package: Package) -> TOMLDocument:
+    def create_legacy_pyproject_from_package(cls, package: Package) -> TOMLDocument:
         import tomlkit
 
         from poetry.utils.dependency_specification import dependency_to_specification

tests/console/commands/self/test_add_plugins.py

@@ -30,8 +30,9 @@ def assert_plugin_add_result(
     constraint: str,
 ) -> None:
     assert tester.io.fetch_output() == expected
-    dependencies: list[str] = get_self_command_dependencies()
+    dependencies: list[str] | None = get_self_command_dependencies()
 
+    assert dependencies
     assert "poetry-plugin" in dependencies[0]
     assert constraint in dependencies[0]
 

tests/console/commands/self/test_remove_plugins.py

@@ -32,10 +32,15 @@ def install_plugin(installed: Repository) -> None:
     package = ProjectPackage("poetry-instance", __version__)
     plugin = Package("poetry-plugin", "1.2.3")
 
-    package.add_dependency(
-        Dependency(plugin.name, "^1.2.3", groups=[SelfCommand.ADDITIONAL_PACKAGE_GROUP])
+    content = Factory.create_legacy_pyproject_from_package(package)
+    content["dependency-groups"] = tomlkit.table()
+    content["dependency-groups"][SelfCommand.ADDITIONAL_PACKAGE_GROUP] = tomlkit.array(  # type: ignore[index]
+        "[\n]"
     )
-    content = Factory.create_pyproject_from_package(package)
+    content["dependency-groups"][SelfCommand.ADDITIONAL_PACKAGE_GROUP].append(  # type: ignore[index, union-attr, call-arg]
+        Dependency(plugin.name, "^1.2.3").to_pep_508()
+    )
+
     system_pyproject_file = SelfCommand.get_default_system_pyproject_file()
     with open(system_pyproject_file, "w", encoding="utf-8", newline="") as f:
         f.write(content.as_string())
@@ -64,7 +69,6 @@ def install_plugin(installed: Repository) -> None:
     installed.add_package(plugin)
 
 
-@pytest.mark.xfail(reason="remove command does not support dependency-groups yet")
 def test_remove_installed_package(tester: CommandTester) -> None:
     tester.execute("poetry-plugin")
 
@@ -82,11 +86,9 @@ def test_remove_installed_package(tester: CommandTester) -> None:
 
     dependencies = get_self_command_dependencies()
 
-    assert "poetry-plugin" not in dependencies
     assert not dependencies
 
 
-@pytest.mark.xfail(reason="remove command does not support dependency-groups yet")
 def test_remove_installed_package_dry_run(tester: CommandTester) -> None:
     tester.execute("poetry-plugin --dry-run")
 
@@ -105,4 +107,6 @@ def test_remove_installed_package_dry_run(tester: CommandTester) -> None:
 
     dependencies = get_self_command_dependencies()
 
-    assert "poetry-plugin" in dependencies
+    assert dependencies
+    assert len(dependencies) == 1
+    assert "poetry-plugin" in dependencies[0]

tests/console/commands/self/test_self_command.py

@@ -19,7 +19,7 @@ def example_system_pyproject() -> str:
     package.add_dependency(
         Dependency(plugin.name, "^1.2.3", groups=[SelfCommand.ADDITIONAL_PACKAGE_GROUP])
     )
-    content = Factory.create_pyproject_from_package(package)
+    content = Factory.create_legacy_pyproject_from_package(package)
     return content.as_string().rstrip("\n")
 
 

tests/console/commands/self/utils.py

@@ -8,7 +8,7 @@
 from poetry.factory import Factory
 
 
-def get_self_command_dependencies(locked: bool = True) -> Array:
+def get_self_command_dependencies(locked: bool = True) -> Array | None:
     from poetry.console.commands.self.self_command import SelfCommand
     from poetry.locations import CONFIG_DIR
 
@@ -23,9 +23,10 @@ def get_self_command_dependencies(locked: bool = True) -> Array:
     poetry = Factory().create_poetry(system_pyproject_file.parent, disable_plugins=True)
 
     pyproject: dict[str, Any] = poetry.file.read()
-    content = pyproject["dependency-groups"]
+    content = pyproject.get("dependency-groups", {})
 
-    assert SelfCommand.ADDITIONAL_PACKAGE_GROUP in content
+    if SelfCommand.ADDITIONAL_PACKAGE_GROUP not in content:
+        return None
 
     dependencies = content[SelfCommand.ADDITIONAL_PACKAGE_GROUP]
     assert isinstance(dependencies, Array)

tests/test_factory.py

@@ -164,7 +164,7 @@ def test_create_pyproject_from_package(
     poetry = Factory().create_poetry(fixture_dir(project))
     package = poetry.package
 
-    pyproject: dict[str, Any] = Factory.create_pyproject_from_package(package)
+    pyproject: dict[str, Any] = Factory.create_legacy_pyproject_from_package(package)
 
     result = pyproject["tool"]["poetry"]
     expected = poetry.pyproject.poetry_config