Skip to content

Commit 6349005

Browse files
committed
refactor: let ArtifactCache handle downloading artifacts that are not yet cached
1 parent 333bb78 commit 6349005

5 files changed

Lines changed: 71 additions & 53 deletions

File tree

src/poetry/installation/executor.py

Lines changed: 13 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import contextlib
44
import csv
5+
import functools
56
import itertools
67
import json
78
import threading
@@ -27,7 +28,6 @@
2728
from poetry.puzzle.exceptions import SolverProblemError
2829
from poetry.utils._compat import decode
2930
from poetry.utils.authenticator import Authenticator
30-
from poetry.utils.cache import ArtifactCache
3131
from poetry.utils.env import EnvCommandError
3232
from poetry.utils.helpers import Downloader
3333
from poetry.utils.helpers import get_file_hash
@@ -76,7 +76,7 @@ def __init__(
7676
else:
7777
self._max_workers = 1
7878

79-
self._artifact_cache = ArtifactCache(cache_dir=config.artifacts_cache_directory)
79+
self._artifact_cache = pool.artifact_cache
8080
self._authenticator = Authenticator(
8181
config, self._io, disable_cache=disable_cache, pool_size=self._max_workers
8282
)
@@ -748,23 +748,11 @@ def _download(self, operation: Install | Update) -> Path:
748748
def _download_link(self, operation: Install | Update, link: Link) -> Path:
749749
package = operation.package
750750

751-
output_dir = self._artifact_cache.get_cache_directory_for_link(link)
752-
# Try to get cached original package for the link provided
751+
# Get original package for the link provided
752+
download_func = functools.partial(self._download_archive, operation)
753753
original_archive = self._artifact_cache.get_cached_archive_for_link(
754-
link, strict=True
754+
link, strict=True, download_func=download_func
755755
)
756-
if original_archive is None:
757-
# No cached original distributions was found, so we download and prepare it
758-
try:
759-
original_archive = self._download_archive(operation, link)
760-
except BaseException:
761-
cache_directory = self._artifact_cache.get_cache_directory_for_link(
762-
link
763-
)
764-
cached_file = cache_directory.joinpath(link.filename)
765-
cached_file.unlink(missing_ok=True)
766-
767-
raise
768756

769757
# Get potential higher prioritized cached archive, otherwise it will fall back
770758
# to the original archive.
@@ -790,7 +778,7 @@ def _download_link(self, operation: Install | Update, link: Link) -> Path:
790778
)
791779
self._write(operation, message)
792780

793-
archive = self._chef.prepare(archive, output_dir=output_dir)
781+
archive = self._chef.prepare(archive, output_dir=original_archive.parent)
794782

795783
# Use the original archive to provide the correct hash.
796784
self._populate_hashes_dict(original_archive, package)
@@ -815,13 +803,13 @@ def _validate_archive_hash(archive: Path, package: Package) -> str:
815803

816804
return archive_hash
817805

818-
def _download_archive(self, operation: Install | Update, link: Link) -> Path:
819-
archive = (
820-
self._artifact_cache.get_cache_directory_for_link(link) / link.filename
821-
)
822-
archive.parent.mkdir(parents=True, exist_ok=True)
823-
824-
downloader = Downloader(link.url, archive, self._authenticator)
806+
def _download_archive(
807+
self,
808+
operation: Install | Update,
809+
url: str,
810+
dest: Path,
811+
) -> None:
812+
downloader = Downloader(url, dest, self._authenticator)
825813
wheel_size = downloader.total_size
826814

827815
operation_message = self.get_operation_message(operation)
@@ -854,8 +842,6 @@ def _download_archive(self, operation: Install | Update, link: Link) -> Path:
854842
with self._lock:
855843
progress.finish()
856844

857-
return archive
858-
859845
def _should_write_operation(self, operation: Operation) -> bool:
860846
return (
861847
not operation.skipped or self._dry_run or self._verbose or not self._enabled

src/poetry/packages/direct_origin.py

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,14 +76,9 @@ def get_package_from_directory(cls, directory: Path) -> Package:
7676

7777
def get_package_from_url(self, url: str) -> Package:
7878
link = Link(url)
79-
artifact = self._artifact_cache.get_cached_archive_for_link(link, strict=True)
80-
81-
if not artifact:
82-
artifact = (
83-
self._artifact_cache.get_cache_directory_for_link(link) / link.filename
84-
)
85-
artifact.parent.mkdir(parents=True, exist_ok=True)
86-
download_file(url, artifact)
79+
artifact = self._artifact_cache.get_cached_archive_for_link(
80+
link, strict=True, download_func=download_file
81+
)
8782

8883
package = self.get_package_from_file(artifact)
8984
package.files = [

src/poetry/utils/cache.py

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from typing import Any
1313
from typing import Generic
1414
from typing import TypeVar
15+
from typing import overload
1516

1617
from poetry.utils._compat import decode
1718
from poetry.utils._compat import encode
@@ -218,18 +219,49 @@ def get_cache_directory_for_git(
218219

219220
return self._get_directory_from_hash(key_parts)
220221

222+
@overload
223+
def get_cached_archive_for_link(
224+
self,
225+
link: Link,
226+
*,
227+
strict: bool,
228+
env: Env | None = ...,
229+
download_func: Callable[[str, Path], None],
230+
) -> Path: ...
231+
232+
@overload
233+
def get_cached_archive_for_link(
234+
self,
235+
link: Link,
236+
*,
237+
strict: bool,
238+
env: Env | None = ...,
239+
download_func: None = ...,
240+
) -> Path | None: ...
241+
221242
def get_cached_archive_for_link(
222243
self,
223244
link: Link,
224245
*,
225246
strict: bool,
226247
env: Env | None = None,
248+
download_func: Callable[[str, Path], None] | None = None,
227249
) -> Path | None:
228250
cache_dir = self.get_cache_directory_for_link(link)
229251

230-
return self._get_cached_archive(
252+
cached_archive = self._get_cached_archive(
231253
cache_dir, strict=strict, filename=link.filename, env=env
232254
)
255+
if cached_archive is None and strict and download_func is not None:
256+
cache_dir.mkdir(parents=True, exist_ok=True)
257+
cached_archive = cache_dir / link.filename
258+
try:
259+
download_func(link.url, cached_archive)
260+
except BaseException:
261+
cached_archive.unlink(missing_ok=True)
262+
raise
263+
264+
return cached_archive
233265

234266
def get_cached_archive_for_git(
235267
self, url: str, reference: str, subdirectory: str | None, env: Env
@@ -246,8 +278,9 @@ def _get_cached_archive(
246278
filename: str | None = None,
247279
env: Env | None = None,
248280
) -> Path | None:
281+
# implication "not strict -> env must not be None"
249282
assert strict or env is not None
250-
# implication "strict -> filename should not be None"
283+
# implication "strict -> filename must not be None"
251284
assert not strict or filename is not None
252285

253286
archives = self._get_cached_archives(cache_dir)

tests/installation/test_executor.py

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
from cleo.io.buffered_io import BufferedIO
2222
from cleo.io.outputs.output import Verbosity
2323
from poetry.core.packages.package import Package
24-
from poetry.core.packages.utils.link import Link
2524
from poetry.core.packages.utils.utils import path_to_url
2625

2726
from poetry.factory import Factory
@@ -593,11 +592,11 @@ def test_executor_should_delete_incomplete_downloads(
593592
side_effect=Exception("Download error"),
594593
)
595594
mocker.patch(
596-
"poetry.installation.executor.ArtifactCache.get_cached_archive_for_link",
595+
"poetry.utils.cache.ArtifactCache._get_cached_archive",
597596
return_value=None,
598597
)
599598
mocker.patch(
600-
"poetry.installation.executor.ArtifactCache.get_cache_directory_for_link",
599+
"poetry.utils.cache.ArtifactCache.get_cache_directory_for_link",
601600
return_value=tmp_path,
602601
)
603602

@@ -823,7 +822,7 @@ def test_executor_should_write_pep610_url_references_for_wheel_urls(
823822
if is_artifact_cached:
824823
link_cached = fixture_dir("distributions") / "demo-0.1.0-py2.py3-none-any.whl"
825824
mocker.patch(
826-
"poetry.installation.executor.ArtifactCache.get_cached_archive_for_link",
825+
"poetry.utils.cache.ArtifactCache.get_cached_archive_for_link",
827826
return_value=link_cached,
828827
)
829828
download_spy = mocker.spy(Executor, "_download_archive")
@@ -861,9 +860,13 @@ def test_executor_should_write_pep610_url_references_for_wheel_urls(
861860
else:
862861
assert package.source_url is not None
863862
download_spy.assert_called_once_with(
864-
mocker.ANY, operation, Link(package.source_url)
863+
mocker.ANY,
864+
operation,
865+
package.source_url,
866+
dest=mocker.ANY,
865867
)
866-
assert download_spy.spy_return.exists(), "cached file should not be deleted"
868+
dest = download_spy.call_args.args[3]
869+
assert dest.exists(), "cached file should not be deleted"
867870

868871

869872
@pytest.mark.parametrize(
@@ -900,12 +903,12 @@ def test_executor_should_write_pep610_url_references_for_non_wheel_urls(
900903
)
901904
download_spy = mocker.spy(Executor, "_download_archive")
902905

903-
if is_sdist_cached | is_wheel_cached:
906+
if is_sdist_cached or is_wheel_cached:
904907
cached_sdist = fixture_dir("distributions") / "demo-0.1.0.tar.gz"
905908
cached_wheel = fixture_dir("distributions") / "demo-0.1.0-py2.py3-none-any.whl"
906909

907-
def mock_get_cached_archive_for_link_func(
908-
_: Link, *, strict: bool, **__: Any
910+
def mock_get_cached_archive_func(
911+
_cache_dir: Path, *, strict: bool, **__: Any
909912
) -> Path | None:
910913
if is_wheel_cached and not strict:
911914
return cached_wheel
@@ -914,8 +917,8 @@ def mock_get_cached_archive_for_link_func(
914917
return None
915918

916919
mocker.patch(
917-
"poetry.installation.executor.ArtifactCache.get_cached_archive_for_link",
918-
side_effect=mock_get_cached_archive_for_link_func,
920+
"poetry.utils.cache.ArtifactCache._get_cached_archive",
921+
side_effect=mock_get_cached_archive_func,
919922
)
920923

921924
package = Package(
@@ -955,9 +958,10 @@ def mock_get_cached_archive_for_link_func(
955958
if expect_artifact_download:
956959
assert package.source_url is not None
957960
download_spy.assert_called_once_with(
958-
mocker.ANY, operation, Link(package.source_url)
961+
mocker.ANY, operation, package.source_url, dest=mocker.ANY
959962
)
960-
assert download_spy.spy_return.exists(), "cached file should not be deleted"
963+
dest = download_spy.call_args.args[3]
964+
assert dest.exists(), "cached file should not be deleted"
961965
else:
962966
download_spy.assert_not_called()
963967

@@ -978,7 +982,7 @@ def test_executor_should_write_pep610_url_references_for_git(
978982
if is_artifact_cached:
979983
link_cached = fixture_dir("distributions") / "demo-0.1.2-py2.py3-none-any.whl"
980984
mocker.patch(
981-
"poetry.installation.executor.ArtifactCache.get_cached_archive_for_git",
985+
"poetry.utils.cache.ArtifactCache.get_cached_archive_for_git",
982986
return_value=link_cached,
983987
)
984988
clone_spy = mocker.spy(Git, "clone")

tests/packages/test_direct_origin.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def test_direct_origin_does_not_download_url_dependency_when_cached(
4343
)
4444
direct_origin = DirectOrigin(artifact_cache)
4545
url = "https://python-poetry.org/distributions/demo-0.1.0-py2.py3-none-any.whl"
46-
mocker.patch(
46+
download_file = mocker.patch(
4747
"poetry.packages.direct_origin.download_file",
4848
side_effect=Exception("download_file should not be called"),
4949
)
@@ -52,5 +52,5 @@ def test_direct_origin_does_not_download_url_dependency_when_cached(
5252

5353
assert package.name == "demo"
5454
artifact_cache.get_cached_archive_for_link.assert_called_once_with(
55-
Link(url), strict=True
55+
Link(url), strict=True, download_func=download_file
5656
)

0 commit comments

Comments
 (0)