Skip to content

Commit

Permalink
Add: Allow to override the URL when using download_async context manager
Browse files Browse the repository at this point in the history
This will allow to define the URL being displayed in the terminal.
  • Loading branch information
bjoernricks committed Jul 4, 2023
1 parent af1e483 commit 93b2286
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions pontos/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
import httpx

from pontos.errors import PontosError
from pontos.typing import SupportsStr

__all__ = (
"AsyncDownloadProgressIterable",
Expand Down Expand Up @@ -94,7 +95,7 @@ def __init__(
self,
*,
content_iterator: AsyncIterator[T],
url: str,
url: SupportsStr,
length: Optional[int],
):
"""
Expand All @@ -107,7 +108,7 @@ def __init__(
length: Length of the content.
"""
self._content_iterator: AsyncIterator[T] = content_iterator
self._url = url
self._url = str(url)
self._length = None if length is None else int(length)

@property
Expand Down Expand Up @@ -146,6 +147,7 @@ async def download_async(
*,
content_length: Optional[int] = None,
chunk_size: int = DEFAULT_CHUNK_SIZE,
url: Optional[str] = None,
) -> AsyncIterator[AsyncDownloadProgressIterable[bytes]]:
"""
An async context manager that returns an AsyncDownloadProgressIterable.
Expand All @@ -157,6 +159,7 @@ async def download_async(
content_length: Optional length of the content to download. If now
provided it is determined from the response if available.
chunk_size: Download the content in chunks of this size.
url: Use a specific URL. If not set the URL of the response is used.
Returns:
A context manager containing an AsyncDownloadProgressIterable
Expand Down Expand Up @@ -185,7 +188,7 @@ async def download_async(
content_length = response.headers.get("content-length")

yield AsyncDownloadProgressIterable(
url=response.url, # type: ignore
url=url if url else response.url,
content_iterator=response.aiter_bytes(chunk_size=chunk_size),
length=content_length,
)
Expand Down

0 comments on commit 93b2286

Please sign in to comment.