create PartialRequirementDownloadCompleter, and use in wheel, install…

…, and download

src/pip/_internal/commands/download.py

@@ -2,13 +2,12 @@
 
 import logging
 import os
-import shutil
 
 from pip._internal.cli import cmdoptions
 from pip._internal.cli.cmdoptions import make_target_python
 from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.cli.status_codes import SUCCESS
-from pip._internal.network.download import BatchDownloader
+from pip._internal.network.download import PartialRequirementDownloadCompleter
 from pip._internal.req.req_tracker import get_requirement_tracker
 from pip._internal.utils.misc import ensure_dir, normalize_path, write_output
 from pip._internal.utils.temp_dir import TempDirectory
@@ -136,23 +135,12 @@ def run(self, options, args):
             reqs, check_supported_wheels=True
         )
 
-        # Download any requirements which were only partially downloaded with
-        # --use-feature=fast-deps.
-        reqs_to_fully_download = [
-            r for r in requirement_set.requirements.values()
-            if r.needs_more_preparation
-        ]
-        links_from_newly_downloaded = []
-        for partially_downloaded_req in reqs_to_fully_download:
-            assert partially_downloaded_req.link
-            links_from_newly_downloaded.append(partially_downloaded_req.link)
-        # Let's download to a temporary directory.
-        tmpdir = TempDirectory(kind="unpack", globally_managed=True).path
-        batch_download = BatchDownloader(session,
-                                         progress_bar=options.progress_bar)
-        download_iterable = batch_download(links_from_newly_downloaded, tmpdir)
-        for _link, (filepath, _) in download_iterable:
-            shutil.copy(filepath, options.download_dir)
+        # Download any requirements which were only fetched by metadata.
+        download_completer = PartialRequirementDownloadCompleter(
+            session,
+            progress_bar=options.progress_bar,
+            download_dir=options.download_dir)
+        download_completer.complete_requirement_downloads(requirement_set)
 
         downloaded = []  # type: List[str]
         for req in requirement_set.requirements.values():

src/pip/_internal/commands/install.py

@@ -18,6 +18,7 @@
 from pip._internal.cli.status_codes import ERROR, SUCCESS
 from pip._internal.exceptions import CommandError, InstallationError
 from pip._internal.locations import distutils_scheme
+from pip._internal.network.download import PartialRequirementDownloadCompleter
 from pip._internal.operations.check import check_install_conflicts
 from pip._internal.req import install_given_reqs
 from pip._internal.req.req_tracker import get_requirement_tracker
@@ -323,6 +324,15 @@ def run(self, options, args):
                 reqs, check_supported_wheels=not options.target_dir
             )
 
+            # Download any requirements which were only fetched by metadata.
+            # Let's download to a temporary directory.
+            tmpdir = TempDirectory(kind="unpack", globally_managed=True).path
+            download_completer = PartialRequirementDownloadCompleter(
+                session,
+                progress_bar=options.progress_bar,
+                download_dir=tmpdir)
+            download_completer.complete_requirement_downloads(requirement_set)
+
             try:
                 pip_req = requirement_set.get_requirement("pip")
             except KeyError:

src/pip/_internal/commands/wheel.py

@@ -11,6 +11,7 @@
 from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.cli.status_codes import SUCCESS
 from pip._internal.exceptions import CommandError
+from pip._internal.network.download import PartialRequirementDownloadCompleter
 from pip._internal.req.req_tracker import get_requirement_tracker
 from pip._internal.utils.misc import ensure_dir, normalize_path
 from pip._internal.utils.temp_dir import TempDirectory
@@ -156,6 +157,13 @@ def run(self, options, args):
             reqs, check_supported_wheels=True
         )
 
+        # Download any requirements which were only fetched by metadata.
+        download_completer = PartialRequirementDownloadCompleter(
+            session,
+            progress_bar=options.progress_bar,
+            download_dir=options.wheel_dir)
+        download_completer.complete_requirement_downloads(requirement_set)
+
         reqs_to_build = [
             r for r in requirement_set.requirements.values()
             if should_build_for_wheel_command(r)

src/pip/_internal/network/download.py

@@ -24,11 +24,13 @@
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
-    from typing import Iterable, Optional, Tuple
+    from typing import Dict, Iterable, Optional, Tuple
 
     from pip._vendor.requests.models import Response
 
     from pip._internal.models.link import Link
+    from pip._internal.models.req_install import InstallRequirement
+    from pip._internal.models.req_set import RequirementSet
     from pip._internal.network.session import PipSession
 
 logger = logging.getLogger(__name__)
@@ -186,7 +188,7 @@ def __init__(
         self._progress_bar = progress_bar
 
     def __call__(self, links, location):
-        # type: (Iterable[Link], str) -> Iterable[Tuple[str, Tuple[str, str]]]
+        # type: (Iterable[Link], str) -> Iterable[Tuple[Link, Tuple[str, str]]]
         """Download the files given by links into location."""
         for link in links:
             try:
@@ -207,4 +209,42 @@ def __call__(self, links, location):
                 for chunk in chunks:
                     content_file.write(chunk)
             content_type = resp.headers.get('Content-Type', '')
-            yield link.url, (filepath, content_type)
+            yield link, (filepath, content_type)
+
+
+class PartialRequirementDownloadCompleter(object):
+
+    def __init__(
+        self,
+        session,           # type: PipSession
+        progress_bar,      # type: str
+        download_dir,      # type: str
+    ):
+        # type: (...) -> None
+        self._batch_downloader = BatchDownloader(session, progress_bar)
+        self._download_dir = download_dir
+
+    def complete_requirement_downloads(self, req_set):
+        # type: (RequirementSet) -> None
+        """Download any requirements which were only partially downloaded with
+        --use-feature=fast-deps."""
+        reqs_to_fully_download = [
+            r for r in req_set.requirements.values()
+            if r.needs_more_preparation
+        ]
+
+        # Map each link to the requirement that owns it. This allows us to set
+        # `req.local_file_path` on the appropriate requirement after passing
+        # all the links at once into BatchDownloader.
+        links_to_fully_download = {}  # type: Dict[Link, InstallRequirement]
+        for req in reqs_to_fully_download:
+            assert req.link
+            links_to_fully_download[req.link] = req
+
+        batch_download = self._batch_downloader(
+            links_to_fully_download.keys(),
+            self._download_dir)
+        for link, (filepath, _) in batch_download:
+            logger.debug("Downloading link %s to %s", link, filepath)
+            req = links_to_fully_download[link]
+            req.local_file_path = filepath

src/pip/_internal/operations/prepare.py

@@ -26,7 +26,7 @@
     VcsHashUnsupported,
 )
 from pip._internal.models.wheel import Wheel
-from pip._internal.network.download import BatchDownloader, Downloader
+from pip._internal.network.download import Downloader
 from pip._internal.network.lazy_wheel import (
     HTTPRangeRequestUnsupported,
     dist_from_wheel_url,
@@ -325,7 +325,6 @@ def __init__(
         self.req_tracker = req_tracker
         self._session = session
         self._download = Downloader(session, progress_bar)
-        self._batch_download = BatchDownloader(session, progress_bar)
         self.finder = finder
 
         # Where still-packed archives should be written to. If None, they are