From 595c32dd5bb3e9c12abdc75caa154718a2c1b8d2 Mon Sep 17 00:00:00 2001 From: Danny McClanahan <1305167+cosmicexplorer@users.noreply.github.com> Date: Tue, 22 Sep 2020 02:13:16 -0700 Subject: [PATCH] create PartialRequirementDownloadCompleter, and use in wheel, install, and download --- src/pip/_internal/commands/download.py | 26 ++++---------- src/pip/_internal/commands/install.py | 10 ++++++ src/pip/_internal/commands/wheel.py | 8 +++++ src/pip/_internal/network/download.py | 46 +++++++++++++++++++++++-- src/pip/_internal/operations/prepare.py | 3 +- 5 files changed, 69 insertions(+), 24 deletions(-) diff --git a/src/pip/_internal/commands/download.py b/src/pip/_internal/commands/download.py index ac16ac79a03..690bfa41bbc 100644 --- a/src/pip/_internal/commands/download.py +++ b/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(): diff --git a/src/pip/_internal/commands/install.py b/src/pip/_internal/commands/install.py index e41660070a0..597eb0749cd 100644 --- a/src/pip/_internal/commands/install.py +++ b/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: diff --git a/src/pip/_internal/commands/wheel.py b/src/pip/_internal/commands/wheel.py index 0f718566bd0..5f09b100e96 100644 --- a/src/pip/_internal/commands/wheel.py +++ b/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) diff --git a/src/pip/_internal/network/download.py b/src/pip/_internal/network/download.py index 56feaabac10..8bdfc1c2f1c 100644 --- a/src/pip/_internal/network/download.py +++ b/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 diff --git a/src/pip/_internal/operations/prepare.py b/src/pip/_internal/operations/prepare.py index f01a4dcd48a..40363aaf39f 100644 --- a/src/pip/_internal/operations/prepare.py +++ b/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