Fix assertion error when determining installation order.

news/10851.bugfix.rst

@@ -0,0 +1 @@
+Fixed assertion error when determining installation order.

src/pip/_internal/resolution/resolvelib/resolver.py

@@ -185,10 +185,7 @@ def get_installation_order(
             return []
 
         graph = self._result.graph
-        weights = get_topological_weights(
-            graph,
-            expected_node_count=len(self._result.mapping) + 1,
-        )
+        weights = get_topological_weights(graph, set(req_set.requirements.keys()))
 
         sorted_items = sorted(
             req_set.requirements.items(),
@@ -199,7 +196,7 @@ def get_installation_order(
 
 
 def get_topological_weights(
-    graph: "DirectedGraph[Optional[str]]", expected_node_count: int
+    graph: "DirectedGraph[Optional[str]]", requirement_keys: Set[str]
 ) -> Dict[Optional[str], int]:
     """Assign weights to each node based on how "deep" they are.
 
@@ -222,6 +219,9 @@ def get_topological_weights(
     don't get stuck in a cycle.
 
     When assigning weight, the longer path (i.e. larger length) is preferred.
+
+    We are only interested in the weights of packages that are in the
+    requirement_keys.
     """
     path: Set[Optional[str]] = set()
     weights: Dict[Optional[str], int] = {}
@@ -237,6 +237,9 @@ def visit(node: Optional[str]) -> None:
             visit(child)
         path.remove(node)
 
+        if node not in requirement_keys:
+            return
+
         last_known_parent_count = weights.get(node, 0)
         weights[node] = max(last_known_parent_count, len(path))
 
@@ -262,6 +265,8 @@ def visit(node: Optional[str]) -> None:
         # Calculate the weight for the leaves.
         weight = len(graph) - 1
         for leaf in leaves:
+            if leaf not in requirement_keys:
+                continue
             weights[leaf] = weight
         # Remove the leaves from the graph, making it simpler.
         for leaf in leaves:
@@ -271,9 +276,10 @@ def visit(node: Optional[str]) -> None:
     # `None` is guaranteed to be the root node by resolvelib.
     visit(None)
 
-    # Sanity checks
-    assert weights[None] == 0
-    assert len(weights) == expected_node_count
+    # Sanity check: all requirement keys should be in the weights,
+    # and no other keys should be in the weights.
+    difference = set(weights.keys()).difference(requirement_keys)
+    assert not difference, difference
 
     return weights
 

tests/unit/resolution_resolvelib/test_resolver.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Tuple, cast
+from typing import Dict, List, Optional, Set, Tuple, cast
 from unittest import mock
 
 import pytest
@@ -103,7 +103,7 @@ def test_new_resolver_get_installation_order(
 
 
 @pytest.mark.parametrize(
-    "name, edges, expected_weights",
+    "name, edges, requirement_keys, expected_weights",
     [
         (
             # From https://github.com/pypa/pip/pull/8127#discussion_r414564664
@@ -116,7 +116,8 @@ def test_new_resolver_get_installation_order(
                 ("three", "four"),
                 ("four", "five"),
             ],
-            {None: 0, "five": 5, "four": 4, "one": 4, "three": 2, "two": 1},
+            {"one", "two", "three", "four", "five"},
+            {"five": 5, "four": 4, "one": 4, "three": 2, "two": 1},
         ),
         (
             "linear",
@@ -127,7 +128,20 @@ def test_new_resolver_get_installation_order(
                 ("three", "four"),
                 ("four", "five"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND restricted",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+            ],
+            {"one", "three", "five"},
+            {"one": 1, "three": 3, "five": 5},
         ),
         (
             "linear AND root -> two",
@@ -139,7 +153,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 (None, "two"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND root -> three",
@@ -151,7 +166,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 (None, "three"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND root -> four",
@@ -163,7 +179,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 (None, "four"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND root -> five",
@@ -175,7 +192,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 (None, "five"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND one -> four",
@@ -187,7 +205,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 ("one", "four"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND two -> four",
@@ -199,7 +218,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 ("two", "four"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND four -> one (cycle)",
@@ -211,7 +231,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 ("four", "one"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND four -> two (cycle)",
@@ -223,7 +244,8 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 ("four", "two"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
         ),
         (
             "linear AND four -> three (cycle)",
@@ -235,16 +257,44 @@ def test_new_resolver_get_installation_order(
                 ("four", "five"),
                 ("four", "three"),
             ],
-            {None: 0, "one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+            {"one", "two", "three", "four", "five"},
+            {"one": 1, "two": 2, "three": 3, "four": 4, "five": 5},
+        ),
+        (
+            "linear AND four -> three (cycle) AND restricted 1-2-3",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("four", "three"),
+            ],
+            {"one", "two", "three"},
+            {"one": 1, "two": 2, "three": 3},
+        ),
+        (
+            "linear AND four -> three (cycle) AND restricted 4-5",
+            [
+                (None, "one"),
+                ("one", "two"),
+                ("two", "three"),
+                ("three", "four"),
+                ("four", "five"),
+                ("four", "three"),
+            ],
+            {"four", "five"},
+            {"four": 4, "five": 5},
         ),
     ],
 )
 def test_new_resolver_topological_weights(
     name: str,
     edges: List[Tuple[Optional[str], Optional[str]]],
+    requirement_keys: Set[str],
     expected_weights: Dict[Optional[str], int],
 ) -> None:
     graph = _make_graph(edges)
 
-    weights = get_topological_weights(graph, len(expected_weights))
+    weights = get_topological_weights(graph, requirement_keys)
     assert weights == expected_weights