Merge pull request #7384 from chrahunt/refactor/cleanup-prepare-3

Cleanup operations.prepare more

Author: chrahunt

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,27 +156,16 @@ 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
         progress_indicator = DownloadProgressProvider(
             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(

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)