From cea5298805aeb70bc6ee60822a0a457459c5d57e Mon Sep 17 00:00:00 2001 From: Punisheroot <44579963+Punisheroot@users.noreply.github.com> Date: Fri, 31 Jul 2026 12:07:07 +0200 Subject: [PATCH] feat(download): prompt to abort background downloads --- src/manage/install_command.py | 8 ++++++- src/manage/urlutils.py | 22 ++++++++++++++++- tests/test_install_command.py | 28 ++++++++++++++++++++++ tests/test_urlutils.py | 45 +++++++++++++++++++++++++++++++++++ 4 files changed, 101 insertions(+), 2 deletions(-) diff --git a/src/manage/install_command.py b/src/manage/install_command.py index 1795bea..bcfe9d8 100644 --- a/src/manage/install_command.py +++ b/src/manage/install_command.py @@ -137,7 +137,13 @@ def _find_creds(url): return None ensure_tree(dest) - urlretrieve(install["url"], dest, on_progress=on_progress, on_auth_request=_find_creds) + urlretrieve( + install["url"], + dest, + on_progress=on_progress, + on_auth_request=_find_creds, + on_cancel=lambda: cmd.ask_yn("Abort download?"), + ) LOGGER.debug("Downloaded to %s", dest) return dest diff --git a/src/manage/urlutils.py b/src/manage/urlutils.py index ef079bc..edfd54a 100644 --- a/src/manage/urlutils.py +++ b/src/manage/urlutils.py @@ -150,6 +150,7 @@ def __init__(self, url, method="GET", headers={}, outfile=None): self.proxy_settings = _proxy_settings_from_env() self._on_progress = None self._on_auth_request = None + self._on_cancel = None def __str__(self): return sanitise_url(self.url) @@ -167,6 +168,11 @@ def on_auth_request(self, url=None): return self.username, self.password return None + def on_cancel(self): + if self._on_cancel: + return self._on_cancel() + return False + def _bits_urlretrieve(request): from _native import (coinitialize, bits_connect, bits_begin, bits_cancel, @@ -234,6 +240,18 @@ def _bits_urlretrieve(request): request.on_progress(progress) last_progress = progress time.sleep(0.1) + except KeyboardInterrupt: + request.on_progress(None) + if job and request.on_cancel(): + try: + bits_cancel(bits, job) + except OSError: + LOGGER.warn("Failed to cancel background download.") + LOGGER.debug("ERROR:", exc_info=True) + else: + if jobfile.is_file(): + unlink(jobfile) + raise except OSError as ex: if job: bits_cancel(bits, job) @@ -542,7 +560,8 @@ def urlopen(url, method="GET", headers={}, on_progress=None, on_auth_request=Non raise RuntimeError("Unable to download from the internet") -def urlretrieve(url, outfile, method="GET", headers={}, chunksize=64 * 1024, on_progress=None, on_auth_request=None): +def urlretrieve(url, outfile, method="GET", headers={}, chunksize=64 * 1024, + on_progress=None, on_auth_request=None, on_cancel=None): scheme, sep, path = url.partition("://") if not sep: scheme = "file" @@ -574,6 +593,7 @@ def on_progress(_): pass request.chunksize = chunksize request._on_progress = on_progress request._on_auth_request = on_auth_request + request._on_cancel = on_cancel first_error = None diff --git a/tests/test_install_command.py b/tests/test_install_command.py index bbb07b1..7e8991d 100644 --- a/tests/test_install_command.py +++ b/tests/test_install_command.py @@ -5,6 +5,7 @@ from manage import install_command as IC from manage import installs +from manage.commands import BaseCommand from manage.exceptions import NoInstallFoundError from manage.logging import LOGGER @@ -73,6 +74,33 @@ def test_merge_existing_index(tmp_path): ] +def test_download_package_yes_aborts_background_download(tmp_path): + class Cmd: + force = False + bundled_dir = None + source = "https://example.com/index.json" + confirm = False + _ask = BaseCommand._ask + ask_yn = BaseCommand.ask_yn + + def urlretrieve(url, dest, **kwargs): + assert url == "https://example.com/download.zip" + assert kwargs["on_cancel"]() is True + dest.write_bytes(b"download") + + dest = tmp_path / "download.zip" + result = IC.download_package( + Cmd(), + {"url": "https://example.com/download.zip"}, + dest, + {}, + urlretrieve=urlretrieve, + ) + + assert result == dest + assert dest.read_bytes() == b"download" + + def test_merge_existing_index_not_found(tmp_path): existing = tmp_path / "index.json" try: diff --git a/tests/test_urlutils.py b/tests/test_urlutils.py index 7d77e78..df32e45 100644 --- a/tests/test_urlutils.py +++ b/tests/test_urlutils.py @@ -421,6 +421,51 @@ def test_bits_urlretrieve_auth(local_withauth, tmp_path): assert dest.read_bytes() == b"Basic placeholder:placeholder" +@pytest.mark.parametrize("cancel,cancel_error", [ + (False, False), + (True, False), + (True, True), +]) +def test_bits_urlretrieve_keyboard_interrupt( + monkeypatch, tmp_path, cancel, cancel_error +): + bits = object() + job = object() + cancelled = [] + cancel_requested = [] + + monkeypatch.setattr(_native, "coinitialize", lambda: None) + monkeypatch.setattr(_native, "bits_connect", lambda: bits) + monkeypatch.setattr(_native, "bits_begin", lambda *a, **k: job) + monkeypatch.setattr(_native, "bits_serialize_job", lambda *a: b"job-id") + + def bits_get_progress(*args): + raise KeyboardInterrupt() + + def bits_cancel(*args): + cancelled.append(args) + if cancel_error: + raise OSError() + + monkeypatch.setattr(_native, "bits_get_progress", bits_get_progress) + monkeypatch.setattr(_native, "bits_cancel", bits_cancel) + + request = UU._Request("https://example.com/download") + request.outfile = tmp_path / "download.zip" + progress = [] + request._on_progress = progress.append + request._on_cancel = lambda: cancel_requested.append(True) or cancel + jobfile = request.outfile.with_suffix(".job") + + with pytest.raises(KeyboardInterrupt): + UU._bits_urlretrieve(request) + + assert cancel_requested == [True] + assert progress == [None] + assert bool(cancelled) == cancel + assert jobfile.is_file() == (not cancel or cancel_error) + + @pytest.fixture def inject_error(): try: