customizable callback when download completed (#127)

* customizable callback when download completed

* typo

* make done_callback a list

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* immutable sequence

* Iterable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test_done_callback

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix typing in done_callback test

* report exception in done callback

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update parfive/config.py typing

Co-authored-by: Stuart Mumford <stuart@cadair.com>

* test_wrongscheme

* typing

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* typing

* test ValueError

* test value error

* test value error

* test_done_callback_error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update parfive/tests/test_downloader.py

Co-authored-by: Stuart Mumford <stuart@cadair.com>

* Update parfive/tests/test_downloader.py

Co-authored-by: Stuart Mumford <stuart@cadair.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use temp files

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update parfive/tests/test_downloader.py

---------

Co-authored-by: Stuart Mumford <stuart@cadair.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

parfive/config.py

@@ -1,7 +1,7 @@
 import os
 import platform
 import warnings
-from typing import Dict, Union, Callable, Optional
+from typing import Dict, Union, Callable, Iterable, Optional
 
 try:
     from typing import Literal  # Added in Python 3.8
@@ -143,6 +143,14 @@ class SessionConfig:
     """
     env: EnvConfig = field(default_factory=EnvConfig)
 
+    done_callbacks: Iterable[Callable[[str, str, Optional[Exception]], None]] = tuple()
+    """
+    A list of functions to be called when a download is completed.
+
+    The signature of the function to be called is `f(filepath: str, url: str, error: Optional[Exception])`.
+    If successful, error will be None, else the occured exception or asyncio.CancelledError.
+    """
+
     @staticmethod
     def _aiofiles_importable():
         try:

parfive/downloader.py

@@ -618,6 +618,8 @@ async def _get_http(
             await asyncio.gather(*tasks)
             # join() waits till all the items in the queue have been processed
             await downloaded_chunk_queue.join()
+            for callback in self.config.done_callbacks:
+                callback(filepath, url, None)
             return str(filepath)
 
         except (Exception, asyncio.CancelledError) as e:
@@ -633,6 +635,8 @@ async def _get_http(
             # computed the filepath, so we have no file to cleanup
             if filepath is not None:
                 remove_file(filepath)
+            for callback in self.config.done_callbacks:
+                callback(filepath, url, e)
             raise FailedDownload(filepath_partial, url, e)
 
         finally:

parfive/tests/test_config.py

@@ -18,6 +18,7 @@ def test_session_config_defaults():
     assert c.https_proxy is None
     assert c.chunksize == 1024
     assert c.use_aiofiles is False
+    assert len(c.done_callbacks) == 0
 
     assert isinstance(c.headers, dict)
     assert "User-Agent" in c.headers

parfive/tests/test_downloader.py

@@ -2,6 +2,7 @@
 import platform
 import threading
 from pathlib import Path
+from tempfile import gettempdir
 from unittest import mock
 from unittest.mock import patch
 
@@ -322,6 +323,15 @@ def test_notaurl(tmpdir):
     assert isinstance(f.errors[0].exception, aiohttp.ClientConnectionError)
 
 
+def test_wrongscheme(tmpdir):
+    tmpdir = str(tmpdir)
+
+    dl = Downloader(progress=False)
+
+    with pytest.raises(ValueError, match="URL must start with either"):
+        dl.enqueue_file("webcal://notaurl.wibble/file", path=tmpdir)
+
+
 def test_retry(tmpdir, testserver):
     tmpdir = str(tmpdir)
     dl = Downloader()
@@ -348,6 +358,25 @@ def test_empty_retry():
     dl.retry(f)
 
 
+def test_done_callback_error(tmpdir, testserver):
+    tmpdir = str(tmpdir)
+
+    def done_callback(filepath, url, error):
+        if error is not None:
+            (Path(gettempdir()) / "callback.error").touch()
+
+    dl = Downloader(config=SessionConfig(done_callbacks=[done_callback]))
+
+    nn = 5
+    for i in range(nn):
+        dl.enqueue_file(testserver.url, path=tmpdir)
+
+    f = dl.download()
+
+    assert (Path(gettempdir()) / "callback.error").exists()
+    (Path(gettempdir()) / "callback.error").unlink()
+
+
 @skip_windows
 @pytest.mark.allow_hosts(True)
 def test_ftp(tmpdir):
@@ -466,6 +495,26 @@ def test_proxy_passed_as_kwargs_to_get(tmpdir, url, proxy):
     ]
 
 
+def test_done_callback(httpserver, tmpdir):
+    tmpdir = str(tmpdir)
+    httpserver.serve_content(
+        "SIMPLE  = T", headers={"Content-Disposition": "attachment; filename=testfile.fits"}
+    )
+
+    def done_callback(filepath, url, error):
+        (Path(gettempdir()) / "callback.done").touch()
+
+    dl = Downloader(config=SessionConfig(done_callbacks=[done_callback]))
+    dl.enqueue_file(httpserver.url, path=Path(tmpdir), max_splits=None)
+
+    assert dl.queued_downloads == 1
+
+    dl.download()
+
+    assert (Path(gettempdir()) / "callback.done").exists()
+    (Path(gettempdir()) / "callback.done").unlink()
+
+
 class CustomThread(threading.Thread):
     def __init__(self, *args, **kwargs):
         self.result = None