move the partial requirement download completion to the bottom of the prepare_more method

Author: cosmicexplorer

news/8896.trivial

@@ -1 +1 @@
-Move the batched download of lazily-fetched wheel files out of the ``Resolver.resolve()`` codepath, into ``PartialRequirementDownloadCompleter``. Consume this in the ``pip download``, ``pip wheel``, and ``pip install`` tasks.
+Separate the batched *download* of lazily-fetched wheel files from the preparation of regularly-downloaded requirements in ``RequirementPreparer.prepare_linked_requirements_more()``.

src/pip/_internal/network/download.py

@@ -24,13 +24,11 @@
 from pip._internal.utils.typing import MYPY_CHECK_RUNNING
 
 if MYPY_CHECK_RUNNING:
-    from typing import Dict, Iterable, Optional, Tuple
+    from typing import 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__)
@@ -176,7 +174,7 @@ def __call__(self, link, location):
         return filepath, content_type
 
 
-class _BatchDownloader(object):
+class BatchDownloader(object):
 
     def __init__(
         self,
@@ -210,37 +208,3 @@ def __call__(self, links, location):
                     content_file.write(chunk)
             content_type = resp.headers.get('Content-Type', '')
             yield link, (filepath, content_type)
-
-
-def complete_partial_requirement_downloads(
-    session,       # type: PipSession
-    progress_bar,  # type: str
-    req_set,       # type: RequirementSet
-    download_dir,  # type: str
-):
-    # type: (...) -> None
-    """Download any requirements which were only partially downloaded with
-    --use-feature=fast-deps."""
-    batch_downloader = _BatchDownloader(session, progress_bar)
-
-    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 = batch_downloader(
-        links_to_fully_download.keys(),
-        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,10 +26,7 @@
     VcsHashUnsupported,
 )
 from pip._internal.models.wheel import Wheel
-from pip._internal.network.download import (
-    Downloader,
-    complete_partial_requirement_downloads,
-)
+from pip._internal.network.download import BatchDownloader, Downloader
 from pip._internal.network.lazy_wheel import (
     HTTPRangeRequestUnsupported,
     dist_from_wheel_url,
@@ -56,7 +53,6 @@
 
     from pip._internal.index.package_finder import PackageFinder
     from pip._internal.models.link import Link
-    from pip._internal.models.req_set import RequirementSet
     from pip._internal.network.session import PipSession
     from pip._internal.req.req_install import InstallRequirement
     from pip._internal.req.req_tracker import RequirementTracker
@@ -329,6 +325,7 @@ 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
@@ -356,10 +353,6 @@ def __init__(
 
         # Should wheels be downloaded lazily?
         self.use_lazy_wheel = lazy_wheel
-        # TODO: this field is only needed in
-        # .complete_partial_requirements(). When the v1 resolver can be
-        # removed, partial downloads can be completed outside of the resolver.
-        self._progress_bar = progress_bar
 
         # Memoized downloaded files, as mapping of url: (path, mime type)
         self._downloaded = {}  # type: Dict[str, Tuple[str, str]]
@@ -498,16 +491,38 @@ def _fetch_metadata_using_lazy_wheel(self, link):
             logger.debug('%s does not support range requests', url)
             return None
 
-    def complete_partial_requirements(self, req_set):
-        # type: (RequirementSet) -> None
+    def _complete_partial_requirements(
+        self,
+        partially_downloaded_reqs,  # type: Iterable[InstallRequirement]
+        parallel_builds=False,      # type: bool
+    ):
+        # type: (...) -> None
         """Download any requirements which were only fetched by metadata."""
-        download_location = self.wheel_download_dir or self.download_dir
-        complete_partial_requirement_downloads(
-            self._session,
-            self._progress_bar,
-            req_set,
-            download_location,
+        # Download to a temporary directory. These will be copied over as
+        # needed for downstream 'download', 'wheel', and 'install' commands.
+        temp_dir = TempDirectory(kind="unpack", globally_managed=True).path
+
+        # 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 partially_downloaded_reqs:
+            assert req.link
+            links_to_fully_download[req.link] = req
+
+        batch_download = self._batch_download(
+            links_to_fully_download.keys(),
+            temp_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
+
+        # This step is necessary to ensure all lazy wheels are processed
+        # successfully by the 'download', 'wheel', and 'install' commands.
+        for req in partially_downloaded_reqs:
+            self._prepare_linked_requirement(req, parallel_builds)
 
     def prepare_linked_requirement(self, req, parallel_builds=False):
         # type: (InstallRequirement, bool) -> Distribution
@@ -537,11 +552,20 @@ def prepare_linked_requirements_more(self, reqs, parallel_builds=False):
                     req.needs_more_preparation = False
 
         # Prepare requirements we found were already downloaded for some
-        # reason. The other downloads will be completed elsewhere.
+        # reason. The other downloads will be completed separately.
+        partially_downloaded_reqs = []  # type: List[InstallRequirement]
         for req in reqs:
-            if not req.needs_more_preparation:
+            if req.needs_more_preparation:
+                partially_downloaded_reqs.append(req)
+            else:
                 self._prepare_linked_requirement(req, parallel_builds)
 
+        # TODO: separate this part out from RequirementPreparer when the v1
+        # resolver can be removed!
+        self._complete_partial_requirements(
+            partially_downloaded_reqs, parallel_builds=parallel_builds,
+        )
+
     def _prepare_linked_requirement(self, req, parallel_builds):
         # type: (InstallRequirement, bool) -> Distribution
         assert req.link

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

@@ -162,11 +162,6 @@ def resolve(self, root_reqs, check_supported_wheels):
 
         reqs = req_set.all_requirements
         self.factory.preparer.prepare_linked_requirements_more(reqs)
-
-        # TODO: extricate this call from the resolver.resolve() code path once
-        # we can drop the v1 resolver.
-        self.factory.preparer.complete_partial_requirements(req_set)
-
         return req_set
 
     def get_installation_order(self, req_set):