customizable callback when download completed

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] = 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

@@ -608,6 +608,8 @@
             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:
@@ -623,6 +625,8 @@
             # 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, None)
             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

@@ -463,3 +463,23 @@ def test_proxy_passed_as_kwargs_to_get(tmpdir, url, proxy):
             "proxy": 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(*args, **kwargs):
+        Path("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("callback.done").exists()
+    Path("callback.done").unlink()