🐛 fix(extras): scope optional deps to requesting edge (#608)

Running `pipdeptree --extras explicit` in an environment where one
dependency asks for `urllib3[socks]` printed `PySocks` under every
`urllib3` in the tree, including the `requests` and `botocore` branches
that only ask for a plain `urllib3`. 🔁 The reverse view had the mirror
problem: `pipdeptree -r -p PySocks` walked out to every `urllib3`
dependent rather than the lone `selenium` that pulled in the `socks`
extra. Closes #607.

A `PackageDAG` keys by node, so `urllib3` is a single node with one
shared child list, while an extra belongs to the edge that requested it.
Resolving `urllib3[socks]` appended `PySocks` to that shared list, and
since the renderers look children up by node key it leaked onto every
parent. An extra edge is path-dependent and a node-keyed graph has
nowhere to store that, so the fix records on each optional edge the
extra that gates it and applies the gate while walking children.
Forward, the walk drops a scoped edge unless the requiring parent asked
for the extra; in reverse, once a descent passes through an extra only
the dependents that asked for it survive. Extras that belong to the
package rather than one edge, such as `--extras active` or an explicit
`--packages foo[extra]`, carry no gate and show everywhere.

A text render revisits a node once per path that reaches it, so the gate
consults a cached set of the nodes that carry a scoped extra and leaves
every other node on its original allocation-free lookup. Package builds
without scoped extras keep their previous cost.

One deliberate choice worth a look: a package listed at the top level
still shows its extra dependencies, since nothing gates a root and the
edge exists. Only the nested occurrences under a non-requesting parent
lose the spurious child.

Author: gaborbernat

src/pipdeptree/_models/dag.py

@@ -5,6 +5,7 @@
 from collections.abc import Iterator, Mapping
 from enum import Enum, auto
 from fnmatch import fnmatch
+from functools import cached_property
 from itertools import chain
 from typing import TYPE_CHECKING, Literal
 
@@ -130,16 +131,38 @@ def get_node_as_parent(self, node_key: str) -> DistPackage | None:
         except KeyError:
             return None
 
-    def get_children(self, node_key: str) -> list[ReqPackage]:
+    def get_children(self, node_key: str, parent: DistPackage | ReqPackage | None = None) -> list[ReqPackage]:
         """
         Get child nodes for a node by its key.
 
         :param node_key: key of the node to get children of
+        :param parent: the edge the node was reached through; gates edge-scoped extras so ``PySocks`` shows only
+            beneath the ``urllib3[socks]`` that asked for it. ``None`` skips gating, which suits top-level nodes and
+            the graph renderers
         :returns: child nodes
 
         """
         node = self.get_node_as_parent(node_key)
-        return self._obj[node] if node else []
+        if node is None:
+            return []
+        children = self._obj[node]
+        if isinstance(parent, ReqPackage):
+            if node_key in self._scoped_parent_keys:
+                return _gate_requested(children, parent.requested_extras)
+            return children
+        if isinstance(parent, DistPackage) and parent.req is not None and (scoped := parent.req.scoped_extra):
+            return _gate_dependents(children, scoped)
+        return children
+
+    @cached_property
+    def _scoped_parent_keys(self) -> frozenset[str]:
+        # Node keys whose children include an edge-scoped extra; only these pay the forward-gating cost. Rendering
+        # revisits a node once per path, so skipping the per-call allocation for the common ungated node matters.
+        return frozenset(
+            node.key
+            for node, children in self._obj.items()
+            if any(isinstance(c, ReqPackage) and c.scoped_extra for c in children)
+        )
 
     def filter_nodes(  # noqa: C901, PLR0912
         self,
@@ -373,6 +396,28 @@ def reverse(self) -> PackageDAG:  # ty: ignore[invalid-method-override]
         return PackageDAG(dict(forward_dag))
 
 
+def _gate_requested(children: list[ReqPackage], requested_extras: frozenset[str]) -> list[ReqPackage]:
+    """Forward gate: keep an extra-scoped edge only when the requiring edge asked for that extra."""
+    requested = {canonicalize_name(e) for e in requested_extras}
+    return [
+        c
+        for c in children
+        if not isinstance(c, ReqPackage) or c.scoped_extra is None or canonicalize_name(c.scoped_extra) in requested
+    ]
+
+
+def _gate_dependents(children: list[ReqPackage], extra: str) -> list[ReqPackage]:
+    """Reverse gate: having descended into a package through ``extra``, keep only dependents that asked for it."""
+    wanted = canonicalize_name(extra)
+    return [
+        c
+        for c in children
+        if not isinstance(c, DistPackage)
+        or c.req is None
+        or wanted in {canonicalize_name(e) for e in c.req.requested_extras}
+    ]
+
+
 def _expand_requested_extras(
     idx: dict[str, DistPackage], requested_extras: Mapping[str, set[str]] | None
 ) -> dict[str, set[str]]:
@@ -396,46 +441,84 @@ def _resolve_extras(
     requested_extras: dict[str, set[str]] | None = None,
 ) -> None:
     """Add extra/optional dependencies to the DAG in-place."""
-    extras_needed = _seed_extras(pkg_deps, idx, extras)
+    seed, unconditional = _seed_extras(pkg_deps, idx, extras)
     for key, wanted in (requested_extras or {}).items():
-        extras_needed.setdefault(key, set()).update(wanted)
-    processed: dict[str, set[str]] = {}
-    # The same (parent, child, extra) triple can be reached through multiple req.extras propagation
-    # paths across rounds; without dedup it would be appended once per path.
-    seen_edges: set[tuple[str, str, str]] = set()
-    while extras_needed:
-        next_round: dict[str, set[str]] = {}
-        for pkg_key, wanted in extras_needed.items():
-            new_extras = wanted - processed.get(pkg_key, set())
-            if not new_extras:
+        seed.setdefault(key, set()).update(wanted)
+        # ``--packages foo[extra]`` is a direct request, so surface it wherever ``foo`` appears.
+        unconditional.update((key, extra) for extra in wanted)
+    _ExtrasExpander(pkg_deps, idx, unconditional).expand(seed)
+
+
+class _ExtrasExpander:
+    """
+    Append extra/optional dependencies to a node-keyed DAG, breadth-first over ``req.extras`` chains.
+
+    Each appended edge records the ``scoped_extra`` gating it, unless its ``(pkg_key, extra)`` is unconditional
+    (active-satisfied or ``--packages`` requested). The renderer reads that tag to place an optional dependency only
+    beneath the parent that asked for the extra instead of under every occurrence of the provider.
+    """
+
+    def __init__(
+        self,
+        pkg_deps: dict[DistPackage, list[ReqPackage]],
+        idx: dict[str, DistPackage],
+        unconditional: set[tuple[str, str]],
+    ) -> None:
+        self._pkg_deps = pkg_deps
+        self._idx = idx
+        self._unconditional = unconditional
+        # The same (parent, child, extra) triple can be reached through multiple req.extras propagation paths across
+        # rounds; without dedup it would be appended once per path.
+        self._seen_edges: set[tuple[str, str, str]] = set()
+        self._processed: dict[str, set[str]] = {}
+
+    def expand(self, pending: dict[str, set[str]]) -> None:
+        while pending:
+            next_round: dict[str, set[str]] = {}
+            for pkg_key, wanted in pending.items():
+                if new_extras := wanted - self._processed.get(pkg_key, set()):
+                    self._processed.setdefault(pkg_key, set()).update(new_extras)
+                    self._attach(pkg_key, new_extras, next_round)
+            pending = next_round
+
+    def _attach(self, pkg_key: str, new_extras: set[str], next_round: dict[str, set[str]]) -> None:
+        dist_pkg = self._idx.get(pkg_key)
+        if dist_pkg is None or dist_pkg not in self._pkg_deps:
+            return
+        for req, extra_name, dep_key in dist_pkg.requires_for_extras(frozenset(new_extras)):
+            if (dist := self._idx.get(dep_key)) is None:
                 continue
-            processed.setdefault(pkg_key, set()).update(new_extras)
-            dist_pkg = idx.get(pkg_key)
-            if dist_pkg is None or dist_pkg not in pkg_deps:
+            edge_key = (dist_pkg.key, dist.key, extra_name)
+            if edge_key in self._seen_edges:
                 continue
-            for req, extra_name, dep_key in dist_pkg.requires_for_extras(frozenset(new_extras)):
-                if (dist := idx.get(dep_key)) is None:
-                    continue
-                edge_key = (dist_pkg.key, dist.key, extra_name)
-                if edge_key in seen_edges:
-                    continue
-                seen_edges.add(edge_key)
-                req.name = dist.project_name
-                pkg_deps[dist_pkg].append(ReqPackage(req, dist, extra=extra_name))
-                if req.extras:
-                    next_round.setdefault(dist.key, set()).update(req.extras)
-        extras_needed = next_round
+            self._seen_edges.add(edge_key)
+            req.name = dist.project_name
+            unconditional = (pkg_key, extra_name) in self._unconditional
+            scoped_extra = None if unconditional else extra_name
+            self._pkg_deps[dist_pkg].append(ReqPackage(req, dist, extra=extra_name, scoped_extra=scoped_extra))
+            if req.extras:
+                next_round.setdefault(dist.key, set()).update(req.extras)
+                if unconditional:
+                    # An unconditional extra pulls its own sub-extras unconditionally too.
+                    self._unconditional.update((dist.key, sub) for sub in req.extras)
 
 
 def _seed_extras(
     pkg_deps: dict[DistPackage, list[ReqPackage]], idx: dict[str, DistPackage], extras: ExtrasMode
-) -> dict[str, set[str]]:
-    """Collect the extras to resolve: requested extras, plus satisfiable ones in active mode."""
+) -> tuple[dict[str, set[str]], set[tuple[str, str]]]:
+    """
+    Collect the extras to resolve: explicit ``name[extra]`` requests, plus satisfiable ones in active mode.
+
+    Returns the extras keyed by package and the set of ``(pkg_key, extra)`` pairs that are unconditional, i.e.
+    active-satisfied rather than gated behind a specific requiring edge.
+    """
     extras_needed = _collect_explicit_extras(pkg_deps)
+    unconditional: set[tuple[str, str]] = set()
     if extras == "active":
         for pkg_key, satisfied in _collect_satisfied_extras(pkg_deps, idx).items():
             extras_needed.setdefault(pkg_key, set()).update(satisfied)
-    return extras_needed
+            unconditional.update((pkg_key, extra) for extra in satisfied)
+    return extras_needed, unconditional
 
 
 def _collect_explicit_extras(pkg_deps: dict[DistPackage, list[ReqPackage]]) -> dict[str, set[str]]:

src/pipdeptree/_models/package.py

@@ -258,11 +258,22 @@ class ReqPackage(Package):
 
     UNKNOWN_VERSION = "?"
 
-    def __init__(self, obj: Requirement, dist: DistPackage | None = None, extra: str | None = None) -> None:
+    def __init__(
+        self,
+        obj: Requirement,
+        dist: DistPackage | None = None,
+        extra: str | None = None,
+        *,
+        scoped_extra: str | None = None,
+    ) -> None:
         super().__init__(obj.name)
         self._obj = obj
         self.dist = dist
         self.extra = extra
+        # The extra this edge is gated behind, or None to surface it everywhere. Set only for an optional dependency
+        # reached through an explicit ``name[extra]`` request (e.g. ``urllib3[socks]``); extras pulled in by
+        # ``--extras active`` or ``--packages foo[extra]`` stay None. PackageDAG.get_children reads it to gate output.
+        self.scoped_extra = scoped_extra
 
     def render_as_root(self, *, frozen: bool) -> str:
         if not frozen:
@@ -287,6 +298,11 @@ def version_spec(self) -> str | None:
         specs = sorted(map(str, self._obj.specifier), reverse=True)  # `reverse` makes '>' prior to '<'
         return ",".join(specs) if specs else None
 
+    @property
+    def requested_extras(self) -> frozenset[str]:
+        """Extras the requiring package asked for on this edge (e.g. ``urllib3[socks]`` -> ``{'socks'}``)."""
+        return frozenset(self._obj.extras)
+
     @property
     def edge_label(self) -> str:
         version = self.version_spec or "any"

src/pipdeptree/_render/json_tree.py

@@ -64,7 +64,7 @@ def aux(
 
         d["dependencies"] = [
             aux(c, parent=node, cur_chain=[*cur_chain, c.project_name])
-            for c in tree.get_children(node.key)
+            for c in tree.get_children(node.key, node)
             if c.project_name not in cur_chain
         ]
 

src/pipdeptree/_render/rich_text.py

@@ -70,7 +70,7 @@ def _build_tree(  # noqa: PLR0913
     if depth >= max_depth:
         return
 
-    children = tree.get_children(node.key)
+    children = tree.get_children(node.key, node)
     for child in children:
         if child.project_name in cur_chain:
             continue

src/pipdeptree/_render/text.py

@@ -101,7 +101,7 @@ def aux(  # noqa: PLR0913, PLR0917
 
         result = [node_str]
 
-        children = tree.get_children(node.key)
+        children = tree.get_children(node.key, node)
         children_strings = [
             aux(
                 c,
@@ -151,7 +151,7 @@ def aux(
         result = [node_str]
         children = [
             aux(c, node, indent=indent + 2, cur_chain=[*cur_chain, c.project_name], depth=depth + 1)
-            for c in tree.get_children(node.key)
+            for c in tree.get_children(node.key, node)
             if c.project_name not in cur_chain and depth + 1 <= max_depth
         ]
         result += list(chain.from_iterable(children))

tests/_models/test_dag.py

@@ -306,6 +306,44 @@ def test_dag_extras_annotates_extra_name(make_mock_dist: MockDistMaker) -> None:
     assert all(dep.extra == "signedtoken" for dep in extra_deps)
 
 
+def _build_socks_dag(make_mock_dist: MockDistMaker, *, extras: ExtrasMode = "explicit") -> PackageDAG:
+    pkgs = [
+        make_mock_dist("selenium", "4.35.0", requires=["urllib3[socks]>=2.5.0,<3.0"]),
+        make_mock_dist("requests", "2.33.1", requires=["urllib3>=1.26,<3"]),
+        make_mock_dist("urllib3", "2.7.0", requires=["pysocks ; extra == 'socks'"], provides_extras=["socks"]),
+        make_mock_dist("pysocks", "1.7.1"),
+    ]
+    return PackageDAG.from_pkgs(pkgs, extras=extras)
+
+
+def _edge_into(dag: PackageDAG, parent_key: str, child_key: str) -> ReqPackage:
+    return next(dep for dep in dag.get_children(parent_key) if dep.key == child_key)
+
+
+@pytest.mark.parametrize(
+    ("extras", "via", "shown"),
+    [
+        pytest.param("explicit", "requests", False, id="explicit-non-requesting-parent-hides"),
+        pytest.param("explicit", "selenium", True, id="explicit-requesting-parent-shows"),
+        pytest.param("explicit", None, True, id="explicit-no-parent-shows-all"),
+        pytest.param("active", "requests", True, id="active-shows-under-every-parent"),
+    ],
+)
+def test_dag_edge_scoped_extra_visibility(
+    make_mock_dist: MockDistMaker, extras: ExtrasMode, via: str | None, shown: bool
+) -> None:
+    dag = _build_socks_dag(make_mock_dist, extras=extras)
+    parent = _edge_into(dag, via, "urllib3") if via else None
+    assert ("pysocks" in {c.key for c in dag.get_children("urllib3", parent)}) is shown
+
+
+def test_dag_reverse_edge_scoped_extra_keeps_only_requesting_dependent(make_mock_dist: MockDistMaker) -> None:
+    reversed_dag = _build_socks_dag(make_mock_dist).reverse()
+    urllib3_dist = next(d for d in reversed_dag.get_children("pysocks") if d.key == "urllib3")
+    dependents = {d.key for d in reversed_dag.get_children("urllib3", urllib3_dist)}
+    assert dependents == {"selenium"}
+
+
 def _build_requested_extras_pkgs(make_mock_dist: MockDistMaker) -> list[Distribution]:
     return [
         make_mock_dist(

tests/render/test_text.py

@@ -605,3 +605,34 @@ def test_render_text_with_extras(
     render_text(dag, max_depth=float("inf"), encoding=encoding)
     output = capsys.readouterr().out
     assert "extra: signedtoken" in output
+
+
+def _socks_dag(make_mock_dist: MockDistMaker) -> PackageDAG:
+    pkgs = [
+        make_mock_dist("selenium", "4.35.0", requires=["urllib3[socks]>=2.5.0,<3.0"]),
+        make_mock_dist("requests", "2.33.1", requires=["urllib3>=1.26,<3"]),
+        make_mock_dist("urllib3", "2.7.0", requires=["pysocks ; extra == 'socks'"], provides_extras=["socks"]),
+        make_mock_dist("pysocks", "1.7.1"),
+    ]
+    return PackageDAG.from_pkgs(pkgs, extras="explicit")
+
+
+def test_render_text_edge_scoped_extra_only_under_requesting_parent(
+    capsys: pytest.CaptureFixture[str], make_mock_dist: MockDistMaker
+) -> None:
+    render_text(_socks_dag(make_mock_dist), max_depth=float("inf"), encoding="utf-8", list_all=False)
+    lines = capsys.readouterr().out.splitlines()
+    requests_block = lines[lines.index("requests==2.33.1") : lines.index("selenium==4.35.0")]
+    selenium_block = lines[lines.index("selenium==4.35.0") :]
+    assert not any("pysocks" in line for line in requests_block)
+    assert any("pysocks" in line for line in selenium_block)
+
+
+def test_render_text_reverse_edge_scoped_extra_prunes_unrelated_dependents(
+    capsys: pytest.CaptureFixture[str], make_mock_dist: MockDistMaker
+) -> None:
+    reversed_dag = _socks_dag(make_mock_dist).reverse().filter_nodes(["pysocks"], None)
+    render_text(reversed_dag, max_depth=float("inf"), encoding="utf-8", list_all=False)
+    output = capsys.readouterr().out
+    assert "selenium" in output
+    assert "requests" not in output