Skip to content

Commit

Permalink
move the partial requirement download completion to the bottom of the…
Browse files Browse the repository at this point in the history
… prepare_more method
  • Loading branch information
cosmicexplorer committed Sep 24, 2020
1 parent 15e61d6 commit 34d27a5
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 62 deletions.
40 changes: 2 additions & 38 deletions src/pip/_internal/network/download.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -176,7 +174,7 @@ def __call__(self, link, location):
return filepath, content_type


class _BatchDownloader(object):
class BatchDownloader(object):

def __init__(
self,
Expand Down Expand Up @@ -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
62 changes: 43 additions & 19 deletions src/pip/_internal/operations/prepare.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
5 changes: 0 additions & 5 deletions src/pip/_internal/resolution/resolvelib/resolver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 34d27a5

Please sign in to comment.