More refactoring

Author: ichard26

src/pip/_internal/network/download.py

@@ -41,9 +41,10 @@ def _prepare_download(
     resp: Response,
     link: Link,
     progress_bar: str,
-    total_length: Optional[int],
     range_start: Optional[int] = 0,
 ) -> Iterable[bytes]:
+    total_length = _get_http_response_size(resp)
+
     if link.netloc == PyPI.file_storage_domain:
         url = link.show_url
     else:
@@ -170,21 +171,14 @@ def __init__(
     def __call__(self, link: Link, location: str) -> Tuple[str, str]:
         """Download the file given by link into location."""
         resp = _http_get_download(self._session, link)
-        total_length = _get_http_response_size(resp)
         content_type = resp.headers.get("Content-Type", "")
 
         filename = _get_http_response_filename(resp, link)
         filepath = os.path.join(location, filename)
 
         with open(filepath, "wb") as content_file:
-            bytes_received = self._process_response(
-                resp, link, content_file, 0, total_length
-            )
-            # If possible, check for an incomplete download and attempt resuming.
-            if total_length and bytes_received < total_length:
-                self._attempt_resume(
-                    resp, link, content_file, total_length, bytes_received
-                )
+            bytes_received = self._process_response(resp, link, content_file, 0)
+            self._attempt_resume_if_needed(resp, link, content_file, bytes_received)
 
         return filepath, content_type
 
@@ -194,11 +188,11 @@ def _process_response(
         link: Link,
         content_file: BinaryIO,
         bytes_received: int,
-        total_length: Optional[int],
     ) -> int:
         """Process the response and write the chunks to the file."""
+        total_length = _get_http_response_size(resp)
         chunks = _prepare_download(
-            resp, link, self._progress_bar, total_length, range_start=bytes_received
+            resp, link, self._progress_bar, range_start=bytes_received
         )
         return self._write_chunks_to_file(
             chunks, content_file, allow_partial=bool(total_length)
@@ -223,15 +217,15 @@ def _write_chunks_to_file(
 
         return bytes_received
 
-    def _attempt_resume(
+    def _attempt_resume_if_needed(
         self,
         resp: Response,
         link: Link,
         content_file: BinaryIO,
-        total_length: Optional[int],
         bytes_received: int,
     ) -> None:
-        """Attempt to resume the download if connection was dropped."""
+        """Attempt to resume/retry the download if connection was dropped."""
+        total_length = _get_http_response_size(resp)
         etag_or_last_modified = _get_http_response_etag_or_last_modified(resp)
 
         attempts_left = self._resume_retries
@@ -246,28 +240,30 @@ def _attempt_resume(
             )
 
             try:
+                # Try to resume the download using a HTTP range request.
                 resume_resp = _http_get_download(
                     self._session,
                     link,
                     range_start=bytes_received,
                     if_range=etag_or_last_modified,
                 )
 
-                # If the server responded with 200 (e.g. when the file has been
-                # modified on the server or range requests are not supported by
-                # the server), reset the download to start from the beginning.
-                restart = resume_resp.status_code != HTTPStatus.PARTIAL_CONTENT
-                if restart:
+                # Fallback: if the server responded with 200 (i.e., the file has
+                # since been modified or range requests are unsupported) or any
+                # other unexpected status, restart the download from the beginning.
+                must_restart = resume_resp.status_code != HTTPStatus.PARTIAL_CONTENT
+                if must_restart:
                     bytes_received, total_length, etag_or_last_modified = (
                         self._reset_download_state(resume_resp, content_file)
                     )
 
                 bytes_received += self._process_response(
-                    resume_resp, link, content_file, bytes_received, total_length
+                    resume_resp, link, content_file, bytes_received
                 )
             except (ConnectionError, ReadTimeoutError, OSError):
                 continue
 
+        # No more resume attempts. Raise an error if the download is still incomplete.
         if total_length and bytes_received < total_length:
             os.remove(content_file.name)
             raise IncompleteDownloadError(

tests/unit/test_network_download.py

@@ -10,7 +10,6 @@
 from pip._internal.models.link import Link
 from pip._internal.network.download import (
     Downloader,
-    _get_http_response_size,
     _http_get_download,
     _prepare_download,
     parse_content_disposition,
@@ -91,12 +90,10 @@ def test_prepare_download__log(
     if from_cache:
         resp.from_cache = from_cache
     link = Link(url)
-    total_length = _get_http_response_size(resp)
     _prepare_download(
         resp,
         link,
         progress_bar="on",
-        total_length=total_length,
         range_start=range_start,
     )