From 2eefaddadaef826c05291ee914a5bf9e8a9b3482 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Randy=20D=C3=B6ring?=
 <30527984+radoering@users.noreply.github.com>
Date: Sun, 19 Jan 2025 17:00:01 +0100
Subject: [PATCH] fix `poetry add --optional` locking

---
 src/poetry/console/commands/add.py | 24 ++++++++++++++
 tests/console/commands/test_add.py | 50 +++++++++++++++++++-----------
 2 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/src/poetry/console/commands/add.py b/src/poetry/console/commands/add.py
index 3c2cc38d0f6..92992249c40 100644
--- a/src/poetry/console/commands/add.py
+++ b/src/poetry/console/commands/add.py
@@ -335,6 +335,13 @@ def handle(self) -> int:
                 else:
                     poetry_section[constraint_name] = poetry_constraint
 
+            if optional:
+                extra_name = canonicalize_name(optional)
+                # _in_extras must be set after converting the dependency to PEP 508
+                # and adding it to the project section to avoid a redundant extra marker
+                dependency._in_extras = [extra_name]
+                self._add_dependency_to_extras(dependency, extra_name)
+
         # Refresh the locker
         if project_section:
             assert group == MAIN_GROUP
@@ -409,3 +416,20 @@ def notify_about_existing_packages(self, existing_packages: list[str]) -> None:
         for name in existing_packages:
             self.line(f"  - <c1>{name}</c1>")
         self.line(self._hint_update_packages)
+
+    def _add_dependency_to_extras(
+        self, dependency: Dependency, extra_name: NormalizedName
+    ) -> None:
+        extras = dict(self.poetry.package.extras)
+        extra_deps = []
+        replaced = False
+        for dep in extras.get(extra_name, ()):
+            if dep.name == dependency.name:
+                extra_deps.append(dependency)
+                replaced = True
+            else:
+                extra_deps.append(dep)
+        if not replaced:
+            extra_deps.append(dependency)
+        extras[extra_name] = extra_deps
+        self.poetry.package.extras = extras
diff --git a/tests/console/commands/test_add.py b/tests/console/commands/test_add.py
index 57ad7ecf4e9..01e9a013fed 100644
--- a/tests/console/commands/test_add.py
+++ b/tests/console/commands/test_add.py
@@ -822,13 +822,16 @@ def test_add_url_constraint_wheel_with_extras(
     [
         (None, {"my-extra": ["cachy (==0.2.0)"]}),
         (
-            {"other": ["foo>2"]},
-            {"other": ["foo>2"], "my-extra": ["cachy (==0.2.0)"]},
+            {"other": ["tomlkit (<2)"]},
+            {"other": ["tomlkit (<2)"], "my-extra": ["cachy (==0.2.0)"]},
         ),
-        ({"my-extra": ["foo>2"]}, {"my-extra": ["foo>2", "cachy (==0.2.0)"]}),
         (
-            {"my-extra": ["foo>2", "cachy (==0.1.0)", "bar>1"]},
-            {"my-extra": ["foo>2", "cachy (==0.2.0)", "bar>1"]},
+            {"my-extra": ["tomlkit (<2)"]},
+            {"my-extra": ["tomlkit (<2)", "cachy (==0.2.0)"]},
+        ),
+        (
+            {"my-extra": ["tomlkit (<2)", "cachy (==0.1.0)", "pendulum (>1)"]},
+            {"my-extra": ["tomlkit (<2)", "cachy (==0.2.0)", "pendulum (>1)"]},
         ),
     ],
 )
@@ -841,29 +844,22 @@ def test_add_constraint_with_optional(
 ) -> None:
     pyproject: dict[str, Any] = app.poetry.file.read()
     if project_dependencies:
-        pyproject["project"]["dependencies"] = ["foo>1"]
+        pyproject["project"]["dependencies"] = ["tomlkit (<1)"]
         if existing_extras:
             pyproject["project"]["optional-dependencies"] = existing_extras
     else:
-        pyproject["tool"]["poetry"]["dependencies"]["foo"] = "^1.0"
+        pyproject["tool"]["poetry"]["dependencies"]["tomlkit"] = "<1"
     pyproject = cast("TOMLDocument", pyproject)
     app.poetry.file.write(pyproject)
+    app.reset_poetry()
 
     tester.execute("cachy=0.2.0 --optional my-extra")
-    expected = """\
 
-Updating dependencies
-Resolving dependencies...
-
-No dependencies to install or update
-
-Writing lock file
-"""
-
-    assert tester.io.fetch_output() == expected
+    assert tester.io.fetch_output().endswith("Writing lock file\n")
     assert isinstance(tester.command, InstallerCommand)
-    assert tester.command.installer.executor.installations_count == 0
+    assert tester.command.installer.executor.installations_count > 0
 
+    # check pyproject content
     pyproject2: dict[str, Any] = app.poetry.file.read()
     project_content = pyproject2["project"]
     poetry_content = pyproject2["tool"]["poetry"]
@@ -887,6 +883,24 @@ def test_add_constraint_with_optional(
             in tester.io.fetch_error()
         )
 
+    # check lock content
+    if project_dependencies:
+        lock_data = app.poetry.locker.lock_data
+
+        extras = lock_data["extras"]
+        assert list(extras) == sorted(expected_extras)
+        assert extras["my-extra"] == sorted(
+            e.split(" ")[0] for e in expected_extras["my-extra"]
+        )
+
+        added_package: dict[str, Any] | None = None
+        for package in lock_data["package"]:
+            if package["name"] == "cachy":
+                added_package = package
+                break
+        assert added_package is not None
+        assert added_package.get("markers") == 'extra == "my-extra"'
+
 
 def test_add_constraint_with_optional_not_main_group(
     app: PoetryTestApplication, tester: CommandTester