Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

customizable callback when download completed #127

Merged
merged 34 commits into from
Apr 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
4bdbffd
customizable callback when download completed
dreamflasher Jan 20, 2023
d96989e
typo
dreamflasher Jan 20, 2023
93239df
Merge branch 'main' into feature/done_callback
dreamflasher Mar 14, 2023
acef9cd
Merge branch 'main' into feature/done_callback
Cadair Mar 24, 2023
046478c
Merge branch 'main' into feature/done_callback
dreamflasher May 9, 2023
087905e
make done_callback a list
dreamflasher May 9, 2023
7255773
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 9, 2023
b8d7ab7
immutable sequence
dreamflasher May 9, 2023
faba13f
Iterable
dreamflasher May 9, 2023
3369753
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 9, 2023
822f645
test_done_callback
dreamflasher May 10, 2023
b45c31b
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] May 10, 2023
2955c27
Merge branch 'main' into feature/done_callback
Cadair Nov 2, 2023
e0406ed
fix typing in done_callback test
dreamflasher Nov 2, 2023
f84006a
report exception in done callback
dreamflasher Nov 2, 2023
dc08da2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 2, 2023
bc8d503
Update parfive/config.py typing
dreamflasher Nov 2, 2023
b9d4fc0
test_wrongscheme
dreamflasher Nov 2, 2023
c0eb575
typing
dreamflasher Nov 2, 2023
ad4118d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 2, 2023
5f0ddc1
typing
dreamflasher Nov 2, 2023
3d374ec
test ValueError
dreamflasher Nov 2, 2023
c4af7e0
test value error
dreamflasher Nov 2, 2023
2efabf7
test value error
dreamflasher Nov 2, 2023
cf29398
test_done_callback_error
dreamflasher Nov 2, 2023
00c5e27
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 2, 2023
e322102
Update parfive/tests/test_downloader.py
dreamflasher Nov 2, 2023
d218425
Update parfive/tests/test_downloader.py
dreamflasher Nov 2, 2023
bcbefb1
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 2, 2023
bb66953
use temp files
dreamflasher Nov 2, 2023
fed6658
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Nov 2, 2023
d45c2b0
Merge remote-tracking branch 'origin/main' into feature/done_callback
Cadair Apr 4, 2024
1417dd7
Merge branch 'main' into feature/done_callback
Cadair Apr 4, 2024
ea1ca05
Update parfive/tests/test_downloader.py
Cadair Apr 4, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion parfive/config.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 4 additions & 0 deletions parfive/downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
1 change: 1 addition & 0 deletions parfive/tests/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
49 changes: 49 additions & 0 deletions parfive/tests/test_downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand Down
Loading