Support using progress bars when downloading

Author: jwodder

README.rst

@@ -44,10 +44,19 @@ Installation
 
     python3 -m pip install pypi-simple
 
+``pypi-simple`` can optionally make use of tqdm_.  To install it alongside
+``pypi-simple``, specify the ``tqdm`` extra::
+
+    python3 -m pip install "pypi-simple[tqdm]"
+
+.. _tqdm: https://tqdm.github.io
+
 
 Example
 =======
 
+Get information about a package:
+
 >>> from pypi_simple import PyPISimple
 >>> with PyPISimple() as client:
 ...     requests_page = client.get_project_page('requests')
@@ -64,3 +73,16 @@ Example
 'sdist'
 >>> pkg.digests
 {'sha256': '813202ace4d9301a3c00740c700e012fb9f3f8c73ddcfe02ab558a8df6f175fd'}
+
+Download a package with a tqdm progress bar:
+
+.. code:: python
+
+    from pypi_simple import PyPISimple, tqdm_progress_factory
+
+    with PyPISimple() as client:
+        page = client.get_project_page("pypi-simple")
+        pkg = page.packages[-1]
+        client.download_package(
+            pkg, path=pkg.filename, progress=tqdm_progress_factory(),
+        )

docs/conf.py

@@ -19,7 +19,7 @@
 
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
-    "requests": ("https://requests.readthedocs.io/en/master/", None),
+    "requests": ("https://requests.readthedocs.io/en/latest/", None),
 }
 
 exclude_patterns = ["_build"]

docs/high-level-api.rst

@@ -12,6 +12,12 @@ Constants
 .. autodata:: PYPI_SIMPLE_ENDPOINT
 .. autodata:: SUPPORTED_REPOSITORY_VERSION
 
+Progress Trackers
+-----------------
+.. autoclass:: ProgressTracker()
+    :special-members: __enter__, __exit__
+.. autofunction:: tqdm_progress_factory
+
 Exceptions
 ----------
 .. autoexception:: DigestMismatchError()

docs/index.rst

@@ -34,10 +34,19 @@ Installation
 
     python3 -m pip install pypi-simple
 
+``pypi-simple`` can optionally make use of tqdm_.  To install it alongside
+``pypi-simple``, specify the ``tqdm`` extra::
+
+    python3 -m pip install "pypi-simple[tqdm]"
+
+.. _tqdm: https://tqdm.github.io
+
 
 Example
 =======
 
+Get information about a package:
+
 >>> from pypi_simple import PyPISimple
 >>> with PyPISimple() as client:
 ...     requests_page = client.get_project_page('requests')
@@ -55,6 +64,19 @@ Example
 >>> pkg.digests
 {'sha256': '813202ace4d9301a3c00740c700e012fb9f3f8c73ddcfe02ab558a8df6f175fd'}
 
+Download a package with a tqdm progress bar:
+
+.. code:: python
+
+    from pypi_simple import PyPISimple, tqdm_progress_factory
+
+    with PyPISimple() as client:
+        page = client.get_project_page("pypi-simple")
+        pkg = page.packages[-1]
+        client.download_package(
+            pkg, path=pkg.filename, progress=tqdm_progress_factory(),
+        )
+
 
 Indices and tables
 ==================

setup.cfg

@@ -51,6 +51,11 @@ install_requires =
     mailbits ~= 0.2
     packaging >= 20
     requests ~= 2.20
+    typing_extensions; python_version < "3.8"
+
+[options.extras_require]
+tqdm =
+    tqdm
 
 [options.packages.find]
 where = src

src/pypi_simple/__init__.py

@@ -40,6 +40,7 @@
     parse_repo_project_response,
 )
 from .parse_stream import parse_links_stream, parse_links_stream_response
+from .progress import ProgressTracker, tqdm_progress_factory
 from .util import (
     DigestMismatchError,
     NoDigestsError,
@@ -55,6 +56,7 @@
     "Link",
     "NoDigestsError",
     "PYPI_SIMPLE_ENDPOINT",
+    "ProgressTracker",
     "ProjectPage",
     "PyPISimple",
     "SUPPORTED_REPOSITORY_VERSION",
@@ -74,4 +76,5 @@
     "parse_repo_project_page",
     "parse_repo_project_response",
     "parse_simple_index",
+    "tqdm_progress_factory",
 ]

src/pypi_simple/client.py

@@ -1,14 +1,16 @@
 import os
 from pathlib import Path
 import platform
-from typing import Any, AnyStr, Iterator, List, Optional, Tuple, Union
+from types import TracebackType
+from typing import Any, AnyStr, Callable, Iterator, List, Optional, Tuple, Type, Union
 from warnings import warn
 from packaging.utils import canonicalize_name as normalize
 import requests
 from . import PYPI_SIMPLE_ENDPOINT, __url__, __version__
 from .classes import DistributionPackage, IndexPage, ProjectPage
 from .parse_repo import parse_repo_index_response, parse_repo_project_response
 from .parse_stream import parse_links_stream_response
+from .progress import ProgressTracker, null_progress_tracker
 from .util import AbstractDigestChecker, DigestChecker, NullDigestChecker
 
 #: The User-Agent header used for requests; not used when the user provides eir
@@ -96,7 +98,12 @@ def __init__(
     def __enter__(self) -> "PyPISimple":
         return self
 
-    def __exit__(self, _exc_type: Any, _exc_value: Any, _exc_tb: Any) -> None:
+    def __exit__(
+        self,
+        _exc_type: Optional[Type[BaseException]],
+        _exc_val: Optional[BaseException],
+        _exc_tb: Optional[TracebackType],
+    ) -> None:
         self.s.close()
 
     def get_index_page(
@@ -215,6 +222,7 @@ def download_package(
         path: Union[AnyStr, "os.PathLike[AnyStr]"],
         verify: bool = True,
         keep_on_error: bool = False,
+        progress: Optional[Callable[[Optional[int]], ProgressTracker]] = None,
         timeout: Union[float, Tuple[float, float], None] = None,
     ) -> None:
         """
@@ -225,6 +233,13 @@ def download_package(
         If an error occurs while downloading or verifying digests, and
         ``keep_on_error`` is not true, the downloaded file is not saved.
 
+        Download progress can be tracked (e.g., for display by a progress bar)
+        by passing an appropriate callable as the ``progress`` argument.  This
+        callable will be passed the length of the downloaded file, if known,
+        and it must return a `ProgressTracker` — a context manager with an
+        ``update(increment: int)`` method that will be passed the size of each
+        downloaded chunk as each chunk is received.
+
         :param DistributionPackage pkg: the distribution package to download
         :param path:
             the path at which to save the downloaded file; any parent
@@ -234,6 +249,7 @@ def download_package(
         :param bool keep_on_error:
             whether to keep (true) or delete (false) the downloaded file if an
             error occurs
+        :param progress: a callable for contructing a progress tracker
         :param timeout: optional timeout to pass to the ``requests`` call
         :type timeout: Union[float, Tuple[float,float], None]
         :raises requests.HTTPError: if the repository responds with an HTTP
@@ -255,10 +271,18 @@ def download_package(
         with self.s.get(pkg.url, stream=True, timeout=timeout) as r:
             r.raise_for_status()
             try:
-                with target.open("wb") as fp:
-                    for chunk in r.iter_content(65535):
-                        fp.write(chunk)
-                        digester.update(chunk)
+                content_length = int(r.headers["Content-Length"])
+            except (ValueError, KeyError):
+                content_length = None
+            if progress is None:
+                progress = null_progress_tracker()
+            try:
+                with progress(content_length) as p:
+                    with target.open("wb") as fp:
+                        for chunk in r.iter_content(65535):
+                            fp.write(chunk)
+                            digester.update(chunk)
+                            p.update(len(chunk))
                 digester.finalize()
             except Exception:
                 if not keep_on_error:

src/pypi_simple/progress.py

@@ -0,0 +1,93 @@
+import sys
+from types import TracebackType
+from typing import Any, Callable, Optional, Type, TypeVar, cast
+
+if sys.version_info[:2] >= (3, 8):
+    from typing import Protocol, runtime_checkable
+else:
+    from typing_extensions import Protocol, runtime_checkable
+
+
+T = TypeVar("T", bound="ProgressTracker")
+
+
+@runtime_checkable
+class ProgressTracker(Protocol):
+    """
+    A `typing.Protocol` for progress trackers.  A progress tracker must be
+    usable as a context manager whose ``__enter__`` method performs startup &
+    returns itself and whose ``__exit__`` method performs shutdown/cleanup.  In
+    addition, a progress tracker must have an ``update(increment: int)`` method
+    that will be called with the size of each downloaded file chunk.
+    """
+
+    def __enter__(self: T) -> T:
+        ...
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        ...
+
+    def update(self, increment: int) -> None:
+        ...
+
+
+class NullProgressTracker:
+    def __enter__(self) -> "NullProgressTracker":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        pass
+
+    def update(self, increment: int) -> None:
+        pass
+
+
+def null_progress_tracker() -> Callable[[Optional[int]], ProgressTracker]:
+    def factory(_content_length: Optional[int]) -> ProgressTracker:
+        return NullProgressTracker()
+
+    return factory
+
+
+def tqdm_progress_factory(**kwargs: Any) -> Callable[[Optional[int]], ProgressTracker]:
+    """
+    A function for displaying a progress bar with tqdm_ during a download.
+    Naturally, using this requires tqdm to be installed alongside
+    ``pypi-simple``.
+
+    Call `tqdm_progress_factory()` with any arguments you wish to pass to the
+    ``tqdm.tqdm`` constructor, and pass the result as the ``progress`` argument
+    to `PyPISimple.download_package()`.
+
+    .. _tqdm: https://tqdm.github.io
+
+    Example:
+
+    .. code:: python
+
+        with PyPISimple() as client:
+            page = client.get_project_page("pypi-simple")
+            pkg = page.packages[-1]
+            client.download_package(
+                pkg,
+                path=pkg.filename,
+                progress=tqdm_progress_factory(desc="Downloading ..."),
+            )
+    """
+
+    from tqdm import tqdm
+
+    def factory(content_length: Optional[int]) -> ProgressTracker:
+        return cast(ProgressTracker, tqdm(total=content_length, **kwargs))
+
+    return factory

src/pypi_simple/util.py

@@ -97,7 +97,7 @@ def update(self, blob: bytes) -> None:
 
     @abstractmethod
     def finalize(self) -> None:
-        pass
+        ...
 
 
 class NullDigestChecker(AbstractDigestChecker):