diff --git a/src/pip/_internal/operations/prepare.py b/src/pip/_internal/operations/prepare.py index 9f419952f70..eb9e32b8109 100644 --- a/src/pip/_internal/operations/prepare.py +++ b/src/pip/_internal/operations/prepare.py @@ -40,7 +40,6 @@ from pip._internal.utils.misc import ( ask_path_exists, backup_dir, - consume, display_path, format_size, hide_url, @@ -58,7 +57,7 @@ if MYPY_CHECK_RUNNING: from typing import ( - Any, Callable, IO, List, Optional, Tuple, + Callable, Iterable, List, Optional, Tuple, ) from mypy_extensions import TypedDict @@ -122,14 +121,12 @@ def _get_http_response_size(resp): return None -def _download_url( +def _prepare_download( resp, # type: Response link, # type: Link - content_file, # type: IO[Any] - hashes, # type: Optional[Hashes] progress_bar # type: str ): - # type: (...) -> None + # type: (...) -> Iterable[bytes] total_length = _get_http_response_size(resp) if link.netloc == PyPI.file_storage_domain: @@ -159,11 +156,6 @@ def _download_url( else: show_progress = False - def written_chunks(chunks): - for chunk in chunks: - content_file.write(chunk) - yield chunk - progress_indicator = _progress_indicator if show_progress: # We don't show progress on cached responses @@ -171,15 +163,9 @@ def written_chunks(chunks): progress_bar, max=total_length ) - downloaded_chunks = written_chunks( - progress_indicator( - response_chunks(resp, CONTENT_CHUNK_SIZE) - ) + return progress_indicator( + response_chunks(resp, CONTENT_CHUNK_SIZE) ) - if hashes: - hashes.check_against_chunks(downloaded_chunks) - else: - consume(downloaded_chunks) def _copy_file(filename, location, link): @@ -215,10 +201,9 @@ def _copy_file(filename, location, link): def unpack_http_url( link, # type: Link location, # type: str - session, # type: PipSession + downloader, # type: Downloader download_dir=None, # type: Optional[str] hashes=None, # type: Optional[Hashes] - progress_bar="on" # type: str ): # type: (...) -> None with TempDirectory(kind="unpack") as temp_dir: @@ -235,7 +220,7 @@ def unpack_http_url( else: # let's download to a tmp dir from_path, content_type = _download_http_url( - link, session, temp_dir.path, hashes, progress_bar + link, downloader, temp_dir.path, hashes ) # unpack the archive to the build dir location. even when only @@ -346,10 +331,9 @@ def unpack_file_url( def unpack_url( link, # type: Link location, # type: str - session, # type: PipSession + downloader, # type: Downloader download_dir=None, # type: Optional[str] hashes=None, # type: Optional[Hashes] - progress_bar="on" # type: str ): # type: (...) -> None """Unpack link. @@ -379,10 +363,9 @@ def unpack_url( unpack_http_url( link, location, - session, + downloader, download_dir, hashes=hashes, - progress_bar=progress_bar ) @@ -464,28 +447,65 @@ def _http_get_download(session, link): return resp +class Download(object): + def __init__( + self, + response, # type: Response + filename, # type: str + chunks, # type: Iterable[bytes] + ): + # type: (...) -> None + self.response = response + self.filename = filename + self.chunks = chunks + + +class Downloader(object): + def __init__( + self, + session, # type: PipSession + progress_bar, # type: str + ): + # type: (...) -> None + self._session = session + self._progress_bar = progress_bar + + def __call__(self, link): + # type: (Link) -> Download + try: + resp = _http_get_download(self._session, link) + except requests.HTTPError as e: + logger.critical( + "HTTP error %s while getting %s", e.response.status_code, link + ) + raise + + return Download( + resp, + _get_http_response_filename(resp, link), + _prepare_download(resp, link, self._progress_bar), + ) + + def _download_http_url( link, # type: Link - session, # type: PipSession + downloader, # type: Downloader temp_dir, # type: str hashes, # type: Optional[Hashes] - progress_bar # type: str ): # type: (...) -> Tuple[str, str] """Download link url into temp_dir using provided session""" - try: - resp = _http_get_download(session, link) - except requests.HTTPError as exc: - logger.critical( - "HTTP error %s while getting %s", exc.response.status_code, link, - ) - raise + download = downloader(link) - filename = _get_http_response_filename(resp, link) - file_path = os.path.join(temp_dir, filename) + file_path = os.path.join(temp_dir, download.filename) with open(file_path, 'wb') as content_file: - _download_url(resp, link, content_file, hashes, progress_bar) - return file_path, resp.headers.get('content-type', '') + for chunk in download.chunks: + content_file.write(chunk) + + if hashes: + hashes.check_against_path(file_path) + + return file_path, download.response.headers.get('content-type', '') def _check_download_dir(link, download_dir, hashes): @@ -538,7 +558,7 @@ def __init__( self.src_dir = src_dir self.build_dir = build_dir self.req_tracker = req_tracker - self.session = session + self.downloader = Downloader(session, progress_bar) self.finder = finder # Where still-packed archives should be written to. If None, they are @@ -559,8 +579,6 @@ def __init__( # be combined if we're willing to have non-wheel archives present in # the wheelhouse output by 'pip wheel'. - self.progress_bar = progress_bar - # Is build isolation allowed? self.build_isolation = build_isolation @@ -662,9 +680,8 @@ def prepare_linked_requirement( try: unpack_url( - link, req.source_dir, self.session, download_dir, + link, req.source_dir, self.downloader, download_dir, hashes=hashes, - progress_bar=self.progress_bar ) except requests.HTTPError as exc: logger.critical( diff --git a/tests/unit/test_operations_prepare.py b/tests/unit/test_operations_prepare.py index d9b3781e9c2..cdcc1400437 100644 --- a/tests/unit/test_operations_prepare.py +++ b/tests/unit/test_operations_prepare.py @@ -13,6 +13,7 @@ from pip._internal.models.link import Link from pip._internal.network.session import PipSession from pip._internal.operations.prepare import ( + Downloader, _copy_source_tree, _download_http_url, parse_content_disposition, @@ -44,6 +45,7 @@ def _fake_session_get(*args, **kwargs): session = Mock() session.get = _fake_session_get + downloader = Downloader(session, progress_bar="on") uri = path_to_url(data.packages.joinpath("simple-1.0.tar.gz")) link = Link(uri) @@ -52,7 +54,7 @@ def _fake_session_get(*args, **kwargs): unpack_http_url( link, temp_dir, - session=session, + downloader=downloader, download_dir=None, ) assert set(os.listdir(temp_dir)) == { @@ -131,6 +133,7 @@ def test_unpack_http_url_bad_downloaded_checksum(mock_unpack_file): response = session.get.return_value = MockResponse(contents) response.headers = {'content-type': 'application/x-tar'} response.url = base_url + downloader = Downloader(session, progress_bar="on") download_dir = mkdtemp() try: @@ -140,7 +143,7 @@ def test_unpack_http_url_bad_downloaded_checksum(mock_unpack_file): unpack_http_url( link, 'location', - session=session, + downloader=downloader, download_dir=download_dir, hashes=Hashes({'sha1': [download_hash.hexdigest()]}) ) @@ -228,15 +231,15 @@ def test_download_http_url__no_directory_traversal(tmpdir): 'content-disposition': 'attachment;filename="../out_dir_file"' } session.get.return_value = resp + downloader = Downloader(session, progress_bar="on") download_dir = tmpdir.joinpath('download') os.mkdir(download_dir) file_path, content_type = _download_http_url( link, - session, + downloader, download_dir, hashes=None, - progress_bar='on', ) # The file should be downloaded to download_dir. actual = os.listdir(download_dir)