Ability to restrict package indexes for specific packages (downstream pipenv issue 4637

src/pip/_internal/index/collector.py

@@ -412,16 +412,19 @@ def __init__(
         self,
         session: PipSession,
         search_scope: SearchScope,
+        index_lookup: Optional[Dict[str, str]] = None,
     ) -> None:
         self.search_scope = search_scope
         self.session = session
+        self.index_lookup = index_lookup if index_lookup else {}
 
     @classmethod
     def create(
         cls,
         session: PipSession,
         options: Values,
         suppress_no_index: bool = False,
+        index_lookup: Optional[Dict[str, str]] = None,
     ) -> "LinkCollector":
         """
         :param session: The Session to use to make requests.
@@ -443,10 +446,10 @@ def create(
             find_links=find_links,
             index_urls=index_urls,
             no_index=options.no_index,
+            index_lookup=index_lookup,
         )
         link_collector = LinkCollector(
-            session=session,
-            search_scope=search_scope,
+            session=session, search_scope=search_scope, index_lookup=index_lookup
         )
         return link_collector
 

src/pip/_internal/models/search_scope.py

@@ -3,7 +3,7 @@
 import os
 import posixpath
 import urllib.parse
-from typing import List
+from typing import Dict, List, Optional
 
 from pip._vendor.packaging.utils import canonicalize_name
 
@@ -20,14 +20,15 @@ class SearchScope:
     Encapsulates the locations that pip is configured to search.
     """
 
-    __slots__ = ["find_links", "index_urls", "no_index"]
+    __slots__ = ["find_links", "index_urls", "no_index", "index_lookup"]
 
     @classmethod
     def create(
         cls,
         find_links: List[str],
         index_urls: List[str],
         no_index: bool,
+        index_lookup: Optional[Dict[str, str]] = None,
     ) -> "SearchScope":
         """
         Create a SearchScope object after normalizing the `find_links`.
@@ -62,17 +63,20 @@ def create(
             find_links=built_find_links,
             index_urls=index_urls,
             no_index=no_index,
+            index_lookup=index_lookup,
         )
 
     def __init__(
         self,
         find_links: List[str],
         index_urls: List[str],
         no_index: bool,
+        index_lookup: Optional[Dict[str, str]] = None,
     ) -> None:
         self.find_links = find_links
         self.index_urls = index_urls
         self.no_index = no_index
+        self.index_lookup = index_lookup if index_lookup else {}
 
     def get_formatted_locations(self) -> str:
         lines = []
@@ -129,4 +133,9 @@ def mkurl_pypi_url(url: str) -> str:
                 loc = loc + "/"
             return loc
 
-        return [mkurl_pypi_url(url) for url in self.index_urls]
+        index_urls = self.index_urls
+        if project_name in self.index_lookup:
+            index_urls = [self.index_lookup[project_name]]
+        elif self.index_urls:
+            index_urls = [self.index_urls[0]]
+        return [mkurl_pypi_url(url) for url in index_urls]

tests/unit/test_search_scope.py

@@ -25,8 +25,8 @@ def test_get_formatted_locations_basic_auth(self) -> None:
         assert "links-user:****@page.domain.com" in result
         assert "links-pass" not in result
 
-    def test_get_index_urls_locations(self) -> None:
-        """Check that the canonical name is on all indexes"""
+    def test_get_index_urls_location_default(self) -> None:
+        """Check that the default index url is used when no index is specified."""
         search_scope = SearchScope(
             find_links=[],
             index_urls=["file://index1/", "file://index2"],
@@ -35,7 +35,17 @@ def test_get_index_urls_locations(self) -> None:
         req = install_req_from_line("Complex_Name")
         assert req.name is not None
         actual = search_scope.get_index_urls_locations(req.name)
-        assert actual == [
-            "file://index1/complex-name/",
-            "file://index2/complex-name/",
-        ]
+        assert actual == ["file://index1/complex-name/"]
+
+    def test_get_index_urls_location_specified(self) -> None:
+        """Check that the specified index url is used."""
+        search_scope = SearchScope(
+            find_links=[],
+            index_urls=["file://index1/", "file://index2"],
+            no_index=False,
+            index_lookup={"Complex_Name": "file://index2"},
+        )
+        req = install_req_from_line("Complex_Name")
+        assert req.name is not None
+        actual = search_scope.get_index_urls_locations(req.name)
+        assert actual == ["file://index2/complex-name/"]