Pass Install Extras to Markers (#9553)

Adds support for conflicting dependencies in extras.

Co-authored-by: Randy Döring <30527984+radoering@users.noreply.github.com>

Author: reesehyde

docs/dependency-specification.md

@@ -572,6 +572,109 @@ pathlib2 = { version = "^2.2", markers = "python_version <= '3.4' or sys_platfor
 {{< /tab >}}
 {{< /tabs >}}
 
+### `extra` environment marker
+
+Poetry populates the `extra` marker with each of the selected extras of the root package.
+For example, consider the following dependency:
+```toml
+[project.optional-dependencies]
+paths = [
+    "pathlib2 (>=2.2,<3.0) ; sys_platform == 'win32'"
+]
+```
+
+`pathlib2` will be installed when you install your package with `--extras paths` on a `win32` machine.
+
+#### Exclusive extras
+
+{{% warning %}}
+The first example will only work completely if you configure Poetry to not re-resolve for installation:
+
+```bash
+poetry config installer.re-resolve false
+```
+
+This is a new feature of Poetry 2.0 that may become the default in a future version of Poetry.
+
+{{% /warning %}}
+
+Keep in mind that all combinations of possible extras available in your project need to be compatible with each other.
+This means that in order to use differing or incompatible versions across different combinations, you need to make your
+extra markers *exclusive*. For example, the following installs PyTorch from one source repository with CPU versions
+when the `cuda` extra is *not* specified, while the other installs from another repository with a separate version set
+for GPUs when the `cuda` extra *is* specified:
+
+```toml
+[project]
+name = "torch-example"
+requires-python = ">=3.10"
+dependencies = [
+    "torch (==2.3.1+cpu) ; extra != 'cuda'",
+]
+
+[project.optional-dependencies]
+cuda = [
+    "torch (==2.3.1+cu118)",
+]
+
+[tool.poetry]
+package-mode = false
+
+[tool.poetry.dependencies]
+torch = [
+    { markers = "extra != 'cuda'", source = "pytorch-cpu"},
+    { markers = "extra == 'cuda'", source = "pytorch-cuda"},
+ ]
+
+[[tool.poetry.source]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
+[[tool.poetry.source]]
+name = "pytorch-cuda"
+url = "https://download.pytorch.org/whl/cu118"
+priority = "explicit"
+```
+
+For the CPU case, we have to specify `"extra != 'cuda'"` because the version specified is not compatible with the
+GPU (`cuda`) version.
+
+This same logic applies when you want either-or extras:
+
+```toml
+[project]
+name = "torch-example"
+requires-python = ">=3.10"
+
+[project.optional-dependencies]
+cpu = [
+    "torch (==2.3.1+cpu)",
+]
+cuda = [
+    "torch (==2.3.1+cu118)",
+]
+
+[tool.poetry]
+package-mode = false
+
+[tool.poetry.dependencies]
+torch = [
+    { markers = "extra == 'cpu' and extra != 'cuda'", source = "pytorch-cpu"},
+    { markers = "extra == 'cuda' and extra != 'cpu'", source = "pytorch-cuda"},
+ ]
+
+[[tool.poetry.source]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+priority = "explicit"
+
+[[tool.poetry.source]]
+name = "pytorch-cuda"
+url = "https://download.pytorch.org/whl/cu118"
+priority = "explicit"
+```
+
 ## Multiple constraints dependencies
 
 Sometimes, one of your dependency may have different version ranges depending

poetry.lock

@@ -305,6 +305,7 @@ def _do_install(self) -> int:
                 self._installed_repository.packages,
                 locked_repository.packages,
                 NullIO(),
+                active_root_extras=self._extras,
             )
             # Everything is resolved at this point, so we no longer need
             # to load deferred dependencies (i.e. VCS, URL and path dependencies)

src/poetry/installation/installer.py

@@ -8,6 +8,7 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from typing import TYPE_CHECKING
+from typing import Any
 from typing import ClassVar
 from typing import cast
 
@@ -17,6 +18,7 @@
 from poetry.core.constraints.version import VersionRange
 from poetry.core.packages.utils.utils import get_python_constraint_from_marker
 from poetry.core.version.markers import AnyMarker
+from poetry.core.version.markers import parse_marker
 from poetry.core.version.markers import union as marker_union
 
 from poetry.mixology.incompatibility import Incompatibility
@@ -115,6 +117,7 @@ def __init__(
         io: IO,
         *,
         locked: list[Package] | None = None,
+        active_root_extras: Collection[NormalizedName] | None = None,
     ) -> None:
         self._package = package
         self._pool = pool
@@ -130,6 +133,9 @@ def __init__(
         self._direct_origin_packages: dict[str, Package] = {}
         self._locked: dict[NormalizedName, list[DependencyPackage]] = defaultdict(list)
         self._use_latest: Collection[NormalizedName] = []
+        self._active_root_extras = (
+            frozenset(active_root_extras) if active_root_extras is not None else None
+        )
 
         self._explicit_sources: dict[str, str] = {}
         for package in locked or []:
@@ -416,21 +422,12 @@ def incompatibilities_for(
                         )
                     ]
 
-        _dependencies = [
-            dep
-            for dep in dependencies
-            if dep.name not in self.UNSAFE_PACKAGES
-            and self._python_constraint.allows_any(dep.python_constraint)
-            and (not self._env or dep.marker.validate(self._env.marker_env))
-        ]
-        dependencies = self._get_dependencies_with_overrides(_dependencies, package)
-
         return [
             Incompatibility(
                 [Term(package.to_dependency(), True), Term(dep, False)],
                 DependencyCauseError(),
             )
-            for dep in dependencies
+            for dep in self._get_dependencies_with_overrides(dependencies, package)
         ]
 
     def complete_package(
@@ -480,7 +477,7 @@ def complete_package(
             package = dependency_package.package
             dependency = dependency_package.dependency
             new_dependency = package.without_features().to_dependency()
-            new_dependency.marker = AnyMarker()
+            new_dependency.marker = dependency.marker
 
             # When adding dependency foo[extra] -> foo, preserve foo's source, if it's
             # specified. This prevents us from trying to get foo from PyPI
@@ -497,8 +494,14 @@ def complete_package(
             if dep.name in self.UNSAFE_PACKAGES:
                 continue
 
-            if self._env and not dep.marker.validate(self._env.marker_env):
-                continue
+            if self._env:
+                marker_values = (
+                    self._marker_values(self._active_root_extras)
+                    if package.is_root()
+                    else self._env.marker_env
+                )
+                if not dep.marker.validate(marker_values):
+                    continue
 
             if not package.is_root() and (
                 (dep.is_optional() and dep.name not in optional_dependencies)
@@ -509,6 +512,24 @@ def complete_package(
             ):
                 continue
 
+            # For normal dependency resolution, we have to make sure that root extras
+            # are represented in the markers. This is required to identify mutually
+            # exclusive markers in cases like 'extra == "foo"' and 'extra != "foo"'.
+            # However, for installation with re-resolving (installer.re-resolve=true,
+            # which results in self._env being not None), this spoils the result
+            # because we have to keep extras so that they are uninstalled
+            # when calculating the operations of the transaction.
+            if self._env is None and package.is_root() and dep.in_extras:
+                # The clone is required for installation with re-resolving
+                # without an existing lock file because the root package is used
+                # once for solving and a second time for re-resolving for installation.
+                dep = dep.clone()
+                dep.marker = dep.marker.intersect(
+                    parse_marker(
+                        " or ".join(f'extra == "{extra}"' for extra in dep.in_extras)
+                    )
+                )
+
             _dependencies.append(dep)
 
         if self._load_deferred:
@@ -545,7 +566,7 @@ def complete_package(
         #   • pypiwin32 (219); sys_platform == "win32" and python_version < "3.6"
         duplicates: dict[str, list[Dependency]] = defaultdict(list)
         for dep in dependencies:
-            duplicates[dep.complete_name].append(dep)
+            duplicates[dep.name].append(dep)
 
         dependencies = []
         for dep_name, deps in duplicates.items():
@@ -556,9 +577,39 @@ def complete_package(
             self.debug(f"<debug>Duplicate dependencies for {dep_name}</debug>")
 
             # For dependency resolution, markers of duplicate dependencies must be
-            # mutually exclusive.
-            active_extras = None if package.is_root() else dependency.extras
-            deps = self._resolve_overlapping_markers(package, deps, active_extras)
+            # mutually exclusive. However, we have to take care about duplicates
+            # with differing extras.
+            duplicates_by_extras: dict[str, list[Dependency]] = defaultdict(list)
+            for dep in deps:
+                duplicates_by_extras[dep.complete_name].append(dep)
+
+            if len(duplicates_by_extras) == 1:
+                active_extras = (
+                    self._active_root_extras if package.is_root() else dependency.extras
+                )
+                deps = self._resolve_overlapping_markers(package, deps, active_extras)
+            else:
+                # There are duplicates with different extras.
+                for complete_dep_name, deps_by_extra in duplicates_by_extras.items():
+                    if len(deps_by_extra) > 1:
+                        duplicates_by_extras[complete_dep_name] = (
+                            self._resolve_overlapping_markers(package, deps, None)
+                        )
+                if all(len(d) == 1 for d in duplicates_by_extras.values()) and all(
+                    d1[0].marker.intersect(d2[0].marker).is_empty()
+                    for d1, d2 in itertools.combinations(
+                        duplicates_by_extras.values(), 2
+                    )
+                ):
+                    # Since all markers are mutually exclusive,
+                    # we can trigger overrides.
+                    deps = list(itertools.chain(*duplicates_by_extras.values()))
+                else:
+                    # Too complicated to handle with overrides,
+                    # fallback to basic handling without overrides.
+                    for d in duplicates_by_extras.values():
+                        dependencies.extend(d)
+                    continue
 
             if len(deps) == 1:
                 self.debug(f"<debug>Merging requirements for {dep_name}</debug>")
@@ -909,3 +960,19 @@ def _resolve_overlapping_markers(
         # dependencies by constraint again. After overlapping markers were
         # resolved, there might be new dependencies with the same constraint.
         return self._merge_dependencies_by_constraint(new_dependencies)
+
+    def _marker_values(
+        self, extras: Collection[NormalizedName] | None = None
+    ) -> dict[str, Any]:
+        """
+        Marker values, from `self._env` if present plus the supplied extras
+
+        :param extras: the values to add to the 'extra' marker value
+        """
+        result = self._env.marker_env.copy() if self._env is not None else {}
+        if extras is not None:
+            assert (
+                "extra" not in result
+            ), "'extra' marker key is already present in environment"
+            result["extra"] = set(extras)
+        return result

src/poetry/puzzle/provider.py

@@ -49,14 +49,21 @@ def __init__(
         installed: list[Package],
         locked: list[Package],
         io: IO,
+        active_root_extras: Collection[NormalizedName] | None = None,
     ) -> None:
         self._package = package
         self._pool = pool
         self._installed_packages = installed
         self._locked_packages = locked
         self._io = io
 
-        self._provider = Provider(self._package, self._pool, self._io, locked=locked)
+        self._provider = Provider(
+            self._package,
+            self._pool,
+            self._io,
+            locked=locked,
+            active_root_extras=active_root_extras,
+        )
         self._overrides: list[dict[Package, dict[str, Dependency]]] = []
 
     @property

src/poetry/puzzle/solver.py

@@ -78,7 +78,7 @@ def calculate_operations(
         else:
             priorities = defaultdict(int)
         relevant_result_packages: set[NormalizedName] = set()
-        uninstalls: set[NormalizedName] = set()
+        pending_extra_uninstalls: list[Package] = []  # list for deterministic order
         for result_package in self._result_packages:
             is_unsolicited_extra = False
             if self._marker_env:
@@ -95,11 +95,12 @@ def calculate_operations(
                 else:
                     continue
             else:
-                relevant_result_packages.add(result_package.name)
                 is_unsolicited_extra = extras is not None and (
                     result_package.optional
                     and result_package.name not in extra_packages
                 )
+                if not is_unsolicited_extra:
+                    relevant_result_packages.add(result_package.name)
 
             installed = False
             for installed_package in self._installed_packages:
@@ -108,9 +109,7 @@ def calculate_operations(
 
                     # Extras that were not requested are always uninstalled.
                     if is_unsolicited_extra:
-                        uninstalls.add(installed_package.name)
-                        if installed_package.name not in system_site_packages:
-                            operations.append(Uninstall(installed_package))
+                        pending_extra_uninstalls.append(installed_package)
 
                     # We have to perform an update if the version or another
                     # attribute of the package has changed (source type, url, ref, ...).
@@ -153,6 +152,13 @@ def calculate_operations(
                     op.skip("Not required")
                 operations.append(op)
 
+        uninstalls: set[NormalizedName] = set()
+        for package in pending_extra_uninstalls:
+            if package.name not in (relevant_result_packages | uninstalls):
+                uninstalls.add(package.name)
+                if package.name not in system_site_packages:
+                    operations.append(Uninstall(package))
+
         if with_uninstalls:
             for current_package in self._current_packages:
                 found = current_package.name in (relevant_result_packages | uninstalls)

src/poetry/puzzle/transaction.py

@@ -0,0 +1,28 @@
+[[package]]
+name = "conflicting-dep"
+version = "1.1.0"
+description = ""
+optional = true
+python-versions = "*"
+files = [ ]
+groups = [ "main" ]
+markers = "extra == \"extra-one\" and extra != \"extra-two\""
+
+[[package]]
+name = "conflicting-dep"
+version = "1.2.0"
+description = ""
+optional = true
+python-versions = "*"
+files = [ ]
+groups = [ "main" ]
+markers = "extra != \"extra-one\" and extra == \"extra-two\""
+
+[extras]
+extra-one = [ "conflicting-dep", "conflicting-dep" ]
+extra-two = [ "conflicting-dep", "conflicting-dep" ]
+
+[metadata]
+lock-version = "2.1"
+python-versions = "*"
+content-hash = "123456789"