gh-364: Prompt to abort background downloads on cancel (#388)

Author: Punisheroot

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
 

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
 

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:

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: