refactor: extract common code from Executor._download_archive and hel…

…pers.download_file into Downloader class

src/poetry/installation/executor.py

@@ -15,7 +15,6 @@
 
 from cleo.io.null_io import NullIO
 from poetry.core.packages.utils.link import Link
-from requests.utils import atomic_open
 
 from poetry.installation.chef import Chef
 from poetry.installation.chef import ChefBuildError
@@ -30,6 +29,7 @@
 from poetry.utils.authenticator import Authenticator
 from poetry.utils.cache import ArtifactCache
 from poetry.utils.env import EnvCommandError
+from poetry.utils.helpers import Downloader
 from poetry.utils.helpers import get_file_hash
 from poetry.utils.helpers import pluralize
 from poetry.utils.helpers import remove_directory
@@ -816,10 +816,14 @@ def _validate_archive_hash(archive: Path, package: Package) -> str:
         return archive_hash
 
     def _download_archive(self, operation: Install | Update, link: Link) -> Path:
-        response = self._authenticator.request(
-            "get", link.url, stream=True, io=self._sections.get(id(operation), self._io)
+        archive = (
+            self._artifact_cache.get_cache_directory_for_link(link) / link.filename
         )
-        wheel_size = response.headers.get("content-length")
+        archive.parent.mkdir(parents=True, exist_ok=True)
+
+        downloader = Downloader(link.url, archive, self._authenticator)
+        wheel_size = downloader.total_size
+
         operation_message = self.get_operation_message(operation)
         message = (
             f"  <fg=blue;options=bold>•</> {operation_message}: <info>Downloading...</>"
@@ -841,23 +845,10 @@ def _download_archive(self, operation: Install | Update, link: Link) -> Path:
                 self._sections[id(operation)].clear()
                 progress.start()
 
-        done = 0
-        archive = (
-            self._artifact_cache.get_cache_directory_for_link(link) / link.filename
-        )
-        archive.parent.mkdir(parents=True, exist_ok=True)
-        with atomic_open(archive) as f:
-            for chunk in response.iter_content(chunk_size=4096):
-                if not chunk:
-                    break
-
-                done += len(chunk)
-
-                if progress:
-                    with self._lock:
-                        progress.set_progress(done)
-
-                f.write(chunk)
+        for fetched_size in downloader.download_with_progress(chunk_size=4096):
+            if progress:
+                with self._lock:
+                    progress.set_progress(fetched_size)
 
         if progress:
             with self._lock:

src/poetry/utils/helpers.py

@@ -12,11 +12,15 @@
 
 from collections.abc import Mapping
 from contextlib import contextmanager
+from contextlib import suppress
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import overload
 
+import requests
+
 from requests.utils import atomic_open
 
 from poetry.utils.constants import REQUESTS_TIMEOUT
@@ -100,43 +104,61 @@ def download_file(
     session: Authenticator | Session | None = None,
     chunk_size: int = 1024,
 ) -> None:
-    import requests
-
     from poetry.puzzle.provider import Indicator
 
-    get = requests.get if not session else session.get
-
-    response = get(url, stream=True, timeout=REQUESTS_TIMEOUT)
-    response.raise_for_status()
+    downloader = Downloader(url, dest, session)
 
     set_indicator = False
     with Indicator.context() as update_context:
         update_context(f"Downloading {url}")
 
-        if "Content-Length" in response.headers:
-            try:
-                total_size = int(response.headers["Content-Length"])
-            except ValueError:
-                total_size = 0
-
+        total_size = downloader.total_size
+        if total_size > 0:
             fetched_size = 0
             last_percent = 0
 
             # if less than 1MB, we simply show that we're downloading
             # but skip the updating
             set_indicator = total_size > 1024 * 1024
 
-        with atomic_open(dest) as f:
-            for chunk in response.iter_content(chunk_size=chunk_size):
+        for fetched_size in downloader.download_with_progress(chunk_size):
+            if set_indicator:
+                percent = (fetched_size * 100) // total_size
+                if percent > last_percent:
+                    last_percent = percent
+                    update_context(f"Downloading {url} {percent:3}%")
+
+
+class Downloader:
+    def __init__(
+        self,
+        url: str,
+        dest: Path,
+        session: Authenticator | Session | None = None,
+    ):
+        self._dest = dest
+
+        get = requests.get if not session else session.get
+
+        self._response = get(url, stream=True, timeout=REQUESTS_TIMEOUT)
+        self._response.raise_for_status()
+
+    @cached_property
+    def total_size(self) -> int:
+        total_size = 0
+        if "Content-Length" in self._response.headers:
+            with suppress(ValueError):
+                total_size = int(self._response.headers["Content-Length"])
+        return total_size
+
+    def download_with_progress(self, chunk_size: int = 1024) -> Iterator[int]:
+        fetched_size = 0
+        with atomic_open(self._dest) as f:
+            for chunk in self._response.iter_content(chunk_size=chunk_size):
                 if chunk:
                     f.write(chunk)
-
-                    if set_indicator:
-                        fetched_size += len(chunk)
-                        percent = (fetched_size * 100) // total_size
-                        if percent > last_percent:
-                            last_percent = percent
-                            update_context(f"Downloading {url} {percent:3}%")
+                    fetched_size += len(chunk)
+                    yield fetched_size
 
 
 def get_package_version_display_string(