Skip to content

Commit 333bb78

Browse files
radoeringradoering
committed
refactor: extract common code from Executor._download_archive and helpers.download_file into Downloader class
1 parent e12ca03 commit 333bb78

2 files changed

Lines changed: 55 additions & 42 deletions

File tree

‎src/poetry/installation/executor.py‎

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
from cleo.io.null_io import NullIO
1717
from poetry.core.packages.utils.link import Link
18-
from requests.utils import atomic_open
1918

2019
from poetry.installation.chef import Chef
2120
from poetry.installation.chef import ChefBuildError
@@ -30,6 +29,7 @@
3029
from poetry.utils.authenticator import Authenticator
3130
from poetry.utils.cache import ArtifactCache
3231
from poetry.utils.env import EnvCommandError
32+
from poetry.utils.helpers import Downloader
3333
from poetry.utils.helpers import get_file_hash
3434
from poetry.utils.helpers import pluralize
3535
from poetry.utils.helpers import remove_directory
@@ -816,10 +816,14 @@ def _validate_archive_hash(archive: Path, package: Package) -> str:
816816
return archive_hash
817817

818818
def _download_archive(self, operation: Install | Update, link: Link) -> Path:
819-
response = self._authenticator.request(
820-
"get", link.url, stream=True, io=self._sections.get(id(operation), self._io)
819+
archive = (
820+
self._artifact_cache.get_cache_directory_for_link(link) / link.filename
821821
)
822-
wheel_size = response.headers.get("content-length")
822+
archive.parent.mkdir(parents=True, exist_ok=True)
823+
824+
downloader = Downloader(link.url, archive, self._authenticator)
825+
wheel_size = downloader.total_size
826+
823827
operation_message = self.get_operation_message(operation)
824828
message = (
825829
f" <fg=blue;options=bold>•</> {operation_message}: <info>Downloading...</>"
@@ -841,23 +845,10 @@ def _download_archive(self, operation: Install | Update, link: Link) -> Path:
841845
self._sections[id(operation)].clear()
842846
progress.start()
843847

844-
done = 0
845-
archive = (
846-
self._artifact_cache.get_cache_directory_for_link(link) / link.filename
847-
)
848-
archive.parent.mkdir(parents=True, exist_ok=True)
849-
with atomic_open(archive) as f:
850-
for chunk in response.iter_content(chunk_size=4096):
851-
if not chunk:
852-
break
853-
854-
done += len(chunk)
855-
856-
if progress:
857-
with self._lock:
858-
progress.set_progress(done)
859-
860-
f.write(chunk)
848+
for fetched_size in downloader.download_with_progress(chunk_size=4096):
849+
if progress:
850+
with self._lock:
851+
progress.set_progress(fetched_size)
861852

862853
if progress:
863854
with self._lock:

‎src/poetry/utils/helpers.py‎

Lines changed: 43 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,15 @@
1212

1313
from collections.abc import Mapping
1414
from contextlib import contextmanager
15+
from contextlib import suppress
16+
from functools import cached_property
1517
from pathlib import Path
1618
from typing import TYPE_CHECKING
1719
from typing import Any
1820
from typing import overload
1921

22+
import requests
23+
2024
from requests.utils import atomic_open
2125

2226
from poetry.utils.constants import REQUESTS_TIMEOUT
@@ -100,43 +104,61 @@ def download_file(
100104
session: Authenticator | Session | None = None,
101105
chunk_size: int = 1024,
102106
) -> None:
103-
import requests
104-
105107
from poetry.puzzle.provider import Indicator
106108

107-
get = requests.get if not session else session.get
108-
109-
response = get(url, stream=True, timeout=REQUESTS_TIMEOUT)
110-
response.raise_for_status()
109+
downloader = Downloader(url, dest, session)
111110

112111
set_indicator = False
113112
with Indicator.context() as update_context:
114113
update_context(f"Downloading {url}")
115114

116-
if "Content-Length" in response.headers:
117-
try:
118-
total_size = int(response.headers["Content-Length"])
119-
except ValueError:
120-
total_size = 0
121-
115+
total_size = downloader.total_size
116+
if total_size > 0:
122117
fetched_size = 0
123118
last_percent = 0
124119

125120
# if less than 1MB, we simply show that we're downloading
126121
# but skip the updating
127122
set_indicator = total_size > 1024 * 1024
128123

129-
with atomic_open(dest) as f:
130-
for chunk in response.iter_content(chunk_size=chunk_size):
124+
for fetched_size in downloader.download_with_progress(chunk_size):
125+
if set_indicator:
126+
percent = (fetched_size * 100) // total_size
127+
if percent > last_percent:
128+
last_percent = percent
129+
update_context(f"Downloading {url} {percent:3}%")
130+
131+
132+
class Downloader:
133+
def __init__(
134+
self,
135+
url: str,
136+
dest: Path,
137+
session: Authenticator | Session | None = None,
138+
):
139+
self._dest = dest
140+
141+
get = requests.get if not session else session.get
142+
143+
self._response = get(url, stream=True, timeout=REQUESTS_TIMEOUT)
144+
self._response.raise_for_status()
145+
146+
@cached_property
147+
def total_size(self) -> int:
148+
total_size = 0
149+
if "Content-Length" in self._response.headers:
150+
with suppress(ValueError):
151+
total_size = int(self._response.headers["Content-Length"])
152+
return total_size
153+
154+
def download_with_progress(self, chunk_size: int = 1024) -> Iterator[int]:
155+
fetched_size = 0
156+
with atomic_open(self._dest) as f:
157+
for chunk in self._response.iter_content(chunk_size=chunk_size):
131158
if chunk:
132159
f.write(chunk)
133-
134-
if set_indicator:
135-
fetched_size += len(chunk)
136-
percent = (fetched_size * 100) // total_size
137-
if percent > last_percent:
138-
last_percent = percent
139-
update_context(f"Downloading {url} {percent:3}%")
160+
fetched_size += len(chunk)
161+
yield fetched_size
140162

141163

142164
def get_package_version_display_string(

0 commit comments

Comments
 (0)