Add support for non-tqdm progress bars on HTTPDownloader (#228)

Allow passing in any object that behaves like a tqdm progress bar instead of just
a boolean. This enables arbitrary custom progress bars or configuring tqdm
objects before passing them to Pooch.

AUTHORS.md

@@ -12,6 +12,7 @@ order by last name) and are considered "The Pooch Developers":
 * [John Leeman](https://github.com/jrleeman)
 * [Daniel McCloy](https://github.com/drammock) - University of Washington, USA (ORCID: [0000-0002-7572-3241](https://orcid.org/0000-0002-7572-3241))
 * [Rémi Rampin](https://github.com/remram44) - New York University, USA (ORCID: [0000-0002-0524-2282](https://www.orcid.org/0000-0002-0524-2282))
+* [Clément Robert](https://github.com/neutrinoceros) - Institut de Planétologie et d'Astrophysique de Grenoble, France (ORCID: [0000-0001-8629-7068](https://orcid.org/0000-0001-8629-7068))
 * [Daniel Shapero](https://github.com/danshapero) - Polar Science Center, University of Washington Applied Physics Lab, USA (ORCID: [0000-0002-3651-0649](https://www.orcid.org/0000-0002-3651-0649))
 * [Santiago Soler](https://github.com/santisoler) - CONICET, Argentina; Instituto Geofísico Sismológico Volponi, Universidad Nacional de San Juan, Argentina (ORCID: [0000-0001-9202-5317](https://www.orcid.org/0000-0001-9202-5317))
 * [Matthew Turk](https://github.com/matthewturk) - University of Illinois at Urbana-Champaign, USA (ORCID: [0000-0002-5294-0198](https://www.orcid.org/0000-0002-5294-0198))

doc/install.rst

@@ -61,7 +61,7 @@ Required:
 Optional:
 
 * `tqdm <https://github.com/tqdm/tqdm>`__: Required to print a download
-  progress bar (see :class:`pooch.HTTPDownloader`).
+  progress bar (see :ref:`tqdm-progressbar` or :ref:`custom-progressbar`).
 * `paramiko <https://github.com/paramiko/paramiko>`__: Required for SFTP
   downloads (see :class:`pooch.SFTPDownloader`).
 

doc/intermediate.rst

@@ -251,8 +251,10 @@ separators** (``'/'``) even on Windows. We will handle the appropriate
 conversions.
 
 
-Printing a download progress bar
---------------------------------
+.. _tqdm-progressbar:
+
+Printing a download progress bar with ``tqdm``
+----------------------------------------------
 
 The :class:`~pooch.HTTPDownloader` can use `tqdm <https://github.com/tqdm/tqdm>`__
 to print a download progress bar. This is turned off by default but can be
@@ -283,3 +285,56 @@ like this:
 
     ``tqdm`` is not installed by default with Pooch. You will have to install
     it separately in order to use this feature.
+
+
+.. _custom-progressbar:
+
+Using custom progress bars
+--------------------------
+
+.. note::
+
+    At the moment, this feature is only available for
+    :class:`pooch.HTTPDownloader`.
+
+Alternatively, you can pass an arbitrary object that behaves like a progress
+that implements the ``update``, ``reset``, and ``close`` methods. ``update``
+should accept a single integer positional argument representing the current
+completion (in bytes), while ``reset`` and ``update`` do not take any argument
+beside ``self``. The object must also have a ``total`` attribute that can be set
+from outside the class.
+In other words, the custom progress bar needs to behave like a ``tqdm`` progress bar.
+Here's a minimal working example of such a custom "progress display" class
+
+.. code:: python
+
+    import sys
+
+    class MinimalProgressDisplay:
+        def __init__(self, total):
+            self.count = 0
+            self.total = total
+
+        def __repr__(self):
+            return str(self.count) + "/" + str(self.total)
+
+        def render(self):
+            print(f"\r{self}", file=sys.stderr, end="")
+
+        def update(self, i):
+            self.count = i
+            self.render()
+
+        def reset(self):
+            self.count = 0
+
+        def close(self):
+            print("", file=sys.stderr)
+
+
+An instance of this class can now be passed to an ``HTTPDownloader`` as
+
+.. code:: python
+
+    pbar = MinimalProgressDisplay(total=None)
+    download = HTTPDownloader(progressbar=pbar)

pooch/downloaders.py

@@ -84,10 +84,11 @@ class HTTPDownloader:  # pylint: disable=too-few-public-methods
 
     Parameters
     ----------
-    progressbar : bool
+    progressbar : bool or an arbitrary progress bar object
         If True, will print a progress bar of the download to standard error
         (stderr). Requires `tqdm <https://github.com/tqdm/tqdm>`__ to be
-        installed.
+        installed. Alternatively, an arbitrary progress bar object can be
+        passed. See :ref:`custom-progressbar` for details.
     chunk_size : int
         Files are streamed *chunk_size* bytes at a time instead of loading
         everything into memory at one. Usually doesn't need to be changed.
@@ -151,7 +152,7 @@ def __init__(self, progressbar=False, chunk_size=1024, **kwargs):
         self.kwargs = kwargs
         self.progressbar = progressbar
         self.chunk_size = chunk_size
-        if self.progressbar and tqdm is None:
+        if self.progressbar is True and tqdm is None:
             raise ValueError("Missing package 'tqdm' required for progress bars.")
 
     def __call__(self, url, output_file, pooch):
@@ -179,8 +180,8 @@ def __call__(self, url, output_file, pooch):
             response = requests.get(url, **kwargs)
             response.raise_for_status()
             content = response.iter_content(chunk_size=self.chunk_size)
-            if self.progressbar:
-                total = int(response.headers.get("content-length", 0))
+            total = int(response.headers.get("content-length", 0))
+            if self.progressbar is True:
                 # Need to use ascii characters on Windows because there isn't
                 # always full unicode support
                 # (see https://github.com/tqdm/tqdm/issues/454)
@@ -193,6 +194,9 @@ def __call__(self, url, output_file, pooch):
                     unit_scale=True,
                     leave=True,
                 )
+            elif self.progressbar:
+                progress = self.progressbar
+                progress.total = total
             for chunk in content:
                 if chunk:
                     output_file.write(chunk)
@@ -245,7 +249,7 @@ class FTPDownloader:  # pylint: disable=too-few-public-methods
     progressbar : bool
         If True, will print a progress bar of the download to standard error
         (stderr). Requires `tqdm <https://github.com/tqdm/tqdm>`__ to be
-        installed.
+        installed. **Custom progress bars are not yet supported.**
     chunk_size : int
         Files are streamed *chunk_size* bytes at a time instead of loading
         everything into memory at one. Usually doesn't need to be changed.
@@ -270,7 +274,7 @@ def __init__(
         self.timeout = timeout
         self.progressbar = progressbar
         self.chunk_size = chunk_size
-        if self.progressbar and tqdm is None:
+        if self.progressbar is True and tqdm is None:
             raise ValueError("Missing package 'tqdm' required for progress bars.")
 
     def __call__(self, url, output_file, pooch):
@@ -352,7 +356,7 @@ class SFTPDownloader:  # pylint: disable=too-few-public-methods
     timeout : int
         Timeout in seconds for sftp socket operations, use None to mean no
         timeout.
-    progressbar : bool
+    progressbar : bool or an arbitrary progress bar object
         If True, will print a progress bar of the download to standard
         error (stderr). Requires `tqdm <https://github.com/tqdm/tqdm>`__ to
         be installed.
@@ -420,6 +424,7 @@ def __call__(self, url, output_file, pooch):
                     unit_scale=True,
                     leave=True,
                 )
+            if self.progressbar:
                 with progress:
 
                     def callback(current, total):

pooch/tests/test_downloaders.py

@@ -164,3 +164,53 @@ def test_downloader_progressbar_sftp(capsys):
         assert printed[:25] == progress
         # Check that the file was actually downloaded
         assert os.path.exists(outfile)
+
+
+def test_downloader_arbitrary_progressbar(capsys):
+    "Setup a downloader function with an arbitrary progress bar class."
+
+    class MinimalProgressDisplay:
+        """A minimalist replacement for tqdm.tqdm"""
+
+        def __init__(self, total):
+            self.count = 0
+            self.total = total
+
+        def __repr__(self):
+            """represent current completion"""
+            return str(self.count) + "/" + str(self.total)
+
+        def render(self):
+            """print self.__repr__ to stderr"""
+            print(f"\r{self}", file=sys.stderr, end="")
+
+        def update(self, i):
+            """modify completion and render"""
+            self.count = i
+            self.render()
+
+        def reset(self):
+            """set counter to 0"""
+            self.count = 0
+
+        @staticmethod
+        def close():
+            """print a new empty line"""
+            print("", file=sys.stderr)
+
+    pbar = MinimalProgressDisplay(total=None)
+    download = HTTPDownloader(progressbar=pbar)
+    with TemporaryDirectory() as local_store:
+        fname = "large-data.txt"
+        url = BASEURL + fname
+        outfile = os.path.join(local_store, "large-data.txt")
+        download(url, outfile, None)
+        # Read stderr and make sure the progress bar is printed only when told
+        captured = capsys.readouterr()
+        printed = captured.err.split("\r")[-1].strip()
+
+        progress = "336/336"
+        assert printed == progress
+
+        # Check that the downloaded file has the right content
+        check_large_data(outfile)