TqdmCallback: add support for custom tqdm subclasses (#1497)

Author: skshetry

fsspec/callbacks.py

@@ -243,7 +243,9 @@ class TqdmCallback(Callback):
     tqdm_kwargs : dict, (optional)
         Any argument accepted by the tqdm constructor.
         See the `tqdm doc <https://tqdm.github.io/docs/tqdm/#__init__>`_.
-        Will be forwarded to tqdm.
+        Will be forwarded to `tqdm_cls`.
+    tqdm_cls: (optional)
+        subclass of `tqdm.tqdm`. If not passed, it will default to `tqdm.tqdm`.
 
     Examples
     --------
@@ -266,6 +268,28 @@ class TqdmCallback(Callback):
             recursive=True,
             callback=TqdmCallback(tqdm_kwargs={"desc": "Your tqdm description"}),
         )
+
+    You can also customize the progress bar by passing a subclass of `tqdm`.
+
+    .. code-block:: python
+
+        class TqdmFormat(tqdm):
+            '''Provides a `total_time` format parameter'''
+            @property
+            def format_dict(self):
+                d = super().format_dict
+                total_time = d["elapsed"] * (d["total"] or 0) / max(d["n"], 1)
+                d.update(total_time=self.format_interval(total_time) + " in total")
+                return d
+
+    >>> with TqdmCallback(
+            tqdm_kwargs={
+                "desc": "desc",
+                "bar_format": "{total_time}: {percentage:.0f}%|{bar}{r_bar}",
+            },
+            tqdm_cls=TqdmFormat,
+        ) as callback:
+            fs.upload(".", path2distant_data, recursive=True, callback=callback)
     """
 
     def __init__(self, tqdm_kwargs=None, *args, **kwargs):
@@ -277,7 +301,7 @@ def __init__(self, tqdm_kwargs=None, *args, **kwargs):
                 "Using TqdmCallback requires tqdm to be installed"
             ) from exce
 
-        self._tqdm_cls = tqdm
+        self._tqdm_cls = kwargs.pop("tqdm_cls", tqdm)
         self._tqdm_kwargs = tqdm_kwargs or {}
         self.tqdm = None
         super().__init__(*args, **kwargs)