refactor: Add _http_get_download() to Downloader and rework methods

Author: ichard26

src/pip/_internal/network/download.py

@@ -8,7 +8,7 @@
 import os
 from collections.abc import Iterable
 from http import HTTPStatus
-from typing import BinaryIO
+from typing import BinaryIO, Mapping
 
 from pip._vendor.requests.models import Response
 from pip._vendor.urllib3.exceptions import ReadTimeoutError
@@ -134,33 +134,9 @@ def _get_http_response_filename(resp: Response, link: Link) -> str:
     return filename
 
 
-def _http_get_download(
-    session: PipSession,
-    link: Link,
-    range_start: int | None = 0,
-    if_range: str | None = None,
-) -> Response:
-    target_url = link.url.split("#", 1)[0]
-    headers = HEADERS.copy()
-    # request a partial download
-    if range_start:
-        headers["Range"] = f"bytes={range_start}-"
-    # make sure the file hasn't changed
-    if if_range:
-        headers["If-Range"] = if_range
-    try:
-        resp = session.get(target_url, headers=headers, stream=True)
-        raise_for_status(resp)
-    except NetworkConnectionError as e:
-        assert e.response is not None
-        logger.critical("HTTP error %s while getting %s", e.response.status_code, link)
-        raise
-    return resp
-
-
 @dataclass
 class _FileDownload:
-    """Stores the state of a single file download."""
+    """Stores the state of a single link download."""
 
     link: Link
     output_file: BinaryIO
@@ -175,7 +151,7 @@ def write_chunk(self, data: bytes) -> None:
         self.bytes_received += len(data)
         self.output_file.write(data)
 
-    def reset_download(self) -> None:
+    def reset_file(self) -> None:
         """Delete any saved data and reset progress to zero."""
         self.output_file.seek(0)
         self.output_file.truncate()
@@ -206,7 +182,7 @@ def batch(
 
     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)
+        resp = self._http_get(link)
         download_size = _get_http_response_size(resp)
 
         filepath = os.path.join(location, _get_http_response_filename(resp, link))
@@ -228,12 +204,6 @@ def _process_response(self, download: _FileDownload, resp: Response) -> None:
             download.size,
             range_start=download.bytes_received,
         )
-        self._write_chunks_to_file(download, chunks)
-
-    def _write_chunks_to_file(
-        self, download: _FileDownload, chunks: Iterable[bytes]
-    ) -> None:
-        """Write the chunks to the file and return the number of bytes received."""
         try:
             for chunk in chunks:
                 download.write_chunk(chunk)
@@ -246,7 +216,6 @@ def _write_chunks_to_file(
 
     def _attempt_resume(self, download: _FileDownload, resp: Response) -> None:
         """Attempt to resume the download if connection was dropped."""
-        etag_or_last_modified = _get_http_response_etag_or_last_modified(resp)
 
         while download.reattempts < self._resume_retries and download.is_incomplete():
             assert download.size is not None
@@ -259,22 +228,14 @@ def _attempt_resume(self, download: _FileDownload, resp: Response) -> None:
             )
 
             try:
-                # Try to resume the download using a HTTP range request.
-                resume_resp = _http_get_download(
-                    self._session,
-                    download.link,
-                    range_start=download.bytes_received,
-                    if_range=etag_or_last_modified,
-                )
-
+                resume_resp = self._http_get_resume(download, should_match=resp)
                 # 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:
-                    download.size, etag_or_last_modified = self._reset_download_state(
-                        download, resume_resp
-                    )
+                    download.reset_file()
+                    download.size = _get_http_response_size(resume_resp)
 
                 self._process_response(download, resume_resp)
             except (ConnectionError, ReadTimeoutError, OSError):
@@ -285,12 +246,27 @@ def _attempt_resume(self, download: _FileDownload, resp: Response) -> None:
             os.remove(download.output_file.name)
             raise IncompleteDownloadError(download)
 
-    def _reset_download_state(
-        self, download: _FileDownload, resp: Response
-    ) -> tuple[int | None, str | None]:
-        """Reset the download state to restart downloading from the beginning."""
-        download.reset_download()
-        total_length = _get_http_response_size(resp)
-        etag_or_last_modified = _get_http_response_etag_or_last_modified(resp)
-
-        return total_length, etag_or_last_modified
+    def _http_get_resume(
+        self, download: _FileDownload, should_match: Response
+    ) -> Response:
+        """Issue a HTTP range request to resume the download."""
+        headers = HEADERS.copy()
+        headers["Range"] = f"bytes={download.bytes_received}-"
+        # If possible, use a conditional range request to avoid corrupted
+        # downloads caused by the remote file changing in-between.
+        if identifier := _get_http_response_etag_or_last_modified(should_match):
+            headers["If-Range"] = identifier
+        return self._http_get(download.link, headers)
+
+    def _http_get(self, link: Link, headers: Mapping[str, str] = HEADERS) -> Response:
+        target_url = link.url_without_fragment
+        try:
+            resp = self._session.get(target_url, headers=headers, stream=True)
+            raise_for_status(resp)
+        except NetworkConnectionError as e:
+            assert e.response is not None
+            logger.critical(
+                "HTTP error %s while getting %s", e.response.status_code, link
+            )
+            raise
+        return resp

tests/unit/test_network_download.py

@@ -12,7 +12,6 @@
 from pip._internal.network.download import (
     Downloader,
     _get_http_response_size,
-    _http_get_download,
     _prepare_download,
     parse_content_disposition,
     sanitize_content_filename,
@@ -149,29 +148,6 @@ def test_sanitize_content_filename__platform_dependent(
     assert sanitize_content_filename(filename) == expected
 
 
-@pytest.mark.parametrize(
-    "range_start, if_range, expected_headers",
-    [
-        (None, None, HEADERS),
-        (1234, None, {**HEADERS, "Range": "bytes=1234-"}),
-        (1234, '"etag"', {**HEADERS, "Range": "bytes=1234-", "If-Range": '"etag"'}),
-    ],
-)
-def test_http_get_download(
-    range_start: int | None,
-    if_range: str | None,
-    expected_headers: dict[str, str],
-) -> None:
-    session = PipSession()
-    session.get = MagicMock()
-    link = Link("http://example.com/foo.tgz")
-    with patch("pip._internal.network.download.raise_for_status"):
-        _http_get_download(session, link, range_start, if_range)
-    session.get.assert_called_once_with(
-        "http://example.com/foo.tgz", headers=expected_headers, stream=True
-    )
-
-
 @pytest.mark.parametrize(
     "content_disposition, default_filename, expected",
     [
@@ -323,7 +299,7 @@ def test_downloader(
     resume_retries: int,
     mock_responses: list[tuple[dict[str, str], int, bytes]],
     # list of (range_start, if_range)
-    expected_resume_args: list[tuple[int | None, int | None]],
+    expected_resume_args: list[tuple[int | None, str | None]],
     # expected_bytes is None means the download should fail
     expected_bytes: bytes | None,
     tmpdir: Path,
@@ -338,9 +314,9 @@ def test_downloader(
         resp.headers = headers
         resp.status_code = status_code
         responses.append(resp)
-    _http_get_download = MagicMock(side_effect=responses)
+    _http_get_mock = MagicMock(side_effect=responses)
 
-    with patch("pip._internal.network.download._http_get_download", _http_get_download):
+    with patch.object(Downloader, "_http_get", _http_get_mock):
         if expected_bytes is None:
             remove = MagicMock(return_value=None)
             with patch("os.remove", remove):
@@ -354,9 +330,12 @@ def test_downloader(
                 downloaded_bytes = downloaded_file.read()
                 assert downloaded_bytes == expected_bytes
 
-    calls = [call(session, link)]  # the initial request
+    calls = [call(link)]  # the initial GET request
     for range_start, if_range in expected_resume_args:
-        calls.append(call(session, link, range_start=range_start, if_range=if_range))
+        headers = {**HEADERS, "Range": f"bytes={range_start}-"}
+        if if_range:
+            headers["If-Range"] = if_range
+        calls.append(call(link, headers))
 
-    # Make sure that the download makes additional requests for resumption
-    _http_get_download.assert_has_calls(calls)
+    # Make sure that the downloader makes additional requests for resumption
+    _http_get_mock.assert_has_calls(calls)