Skip to content

Commit

Permalink
TqdmCallback: add support for custom tqdm subclasses (#1497)
Browse files Browse the repository at this point in the history
  • Loading branch information
skshetry authored Jan 29, 2024
1 parent a408121 commit f21699a
Showing 1 changed file with 26 additions and 2 deletions.
28 changes: 26 additions & 2 deletions fsspec/callbacks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
--------
Expand All @@ -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):
Expand All @@ -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)
Expand Down

0 comments on commit f21699a

Please sign in to comment.