feat!: stream get; pass client; hook guard (#16)

Co-authored-by: JamzumSum <zzzzss990315@gmail.com>

Author: github-actions[bot]

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "AssetsUpdater"
-version = "1.4.1"
+version = "1.4.2"
 description = "Update assets from network."
 authors = ["JamzumSum <zzzzss990315@gmail.com>"]
 license = "MIT"
@@ -25,7 +25,7 @@ rich = { version = "^12.6.0", optional = true }
 optional = false
 
 [tool.poetry.group.test.dependencies]
-pytest = "^7.1.3"
+pytest = "^7.2.0"
 pytest-asyncio = "~0.20.1"
 
 [tool.poetry.group.dev]

pyproject.toml

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+from functools import wraps
 from os import PathLike
 from pathlib import Path
 from typing import Awaitable, Callable, Optional
@@ -13,26 +14,57 @@
 __all__ = ["download"]
 
 
+def _guard(func: Callable):
+    warned = False
+
+    @wraps(func)
+    async def wrapper(*args, **kwds):
+        try:
+            return await func(*args, **kwds)
+        except:
+            nonlocal warned
+            if not warned:
+                log.error("Error caught by guard.", exc_info=True)
+                warned = True
+
+    return wrapper
+
+
 async def download(
     url: URLTypes,
     local: Optional[PathLike] = None,
     buffer: int = 32768,
     echo_progress: Optional[Callable[..., Awaitable]] = None,
-    proxy: ProxiesTypes = ...,
-):
+    *,
+    client: Optional[AsyncClient] = None,
+    proxy: Optional[ProxiesTypes] = None,
+) -> int:
     """async-download a url to a the given path.
 
     :param local: where to download
+    :param buffer: transfer buffer
     :param echo_progress: async function to receive downloaded size and total size as int.
-
+    :param client: use this client, otherwise we will create one and close it on return.
+    :param proxy: download proxy.
     :return: size
+
+    .. warning:: actual proxy = client.proxy | param.proxy
+
+    .. versionchanged:: 1.4.2
+
+        added `client`; `proxy` is set as a keyword-only parameter.
     """
+    if client is None:
+        async with AsyncClient(proxies=proxy) as client:
+            return await download(url, local, buffer, echo_progress, client=client, proxy=proxy)
+
     url = url if isinstance(url, URL) else URL(url)
     remote_name = Path(url.path).name
     local = Path(local or remote_name)
     local.parent.mkdir(parents=True, exist_ok=True)
 
     if local.is_dir():
+        # if local not exist or is a directory
         local = local / remote_name
         log.info(f"Changing local path from {local.parent.as_posix()} to {local.as_posix()}")
         if local.is_dir():
@@ -41,29 +73,29 @@ async def download(
     if local.exists():
         log.warning(f"{local.as_posix()} exists. Overwrite.")
 
-    client_dict = {}
-    if proxy != ...:
-        client_dict["proxies"] = proxy
+    if echo_progress:
+        echo_progress = _guard(echo_progress)
 
     acc = 0
-    async with AsyncClient(**client_dict) as client, aiofiles.open(local, "wb") as f:
-        r = await client.get(url, follow_redirects=True)
-        r.raise_for_status()
-
-        size = int(r.headers.get("Content-Length", -1))
-        it = r.aiter_bytes(chunk_size=buffer)
-
-        log.info(f"Starting download task {url} -> {local}")
-        log.debug(f"File size: {size}")
-
-        async for c in it:
-            if not c:
-                continue
-            if echo_progress:
-                ad, _ = await asyncio.gather(f.write(c), echo_progress(completed=acc, total=size))
-                acc += ad
-            else:
-                acc += await f.write(c)
+    async with aiofiles.open(local, "wb") as f:
+        async with client.stream("GET", url, follow_redirects=True) as r:
+            r.raise_for_status()
+
+            size = int(r.headers.get("Content-Length", -1))
+
+            log.info(f"Starting download task {url} -> {local}")
+            log.debug(f"File size: {size}")
+
+            async for c in r.aiter_bytes(chunk_size=buffer):
+                if not c:
+                    continue
+                if echo_progress:
+                    ad, _ = await asyncio.gather(
+                        f.write(c), echo_progress(completed=acc, total=size)
+                    )
+                    acc += ad
+                else:
+                    acc += await f.write(c)
 
     if size > 0 and acc != size:
         log.error(f"Content-Length is {size} but {acc} downloaded")