Spaces:
Running
Running
| from functools import wraps | |
| class Callback: | |
| """ | |
| Base class and interface for callback mechanism | |
| This class can be used directly for monitoring file transfers by | |
| providing ``callback=Callback(hooks=...)`` (see the ``hooks`` argument, | |
| below), or subclassed for more specialised behaviour. | |
| Parameters | |
| ---------- | |
| size: int (optional) | |
| Nominal quantity for the value that corresponds to a complete | |
| transfer, e.g., total number of tiles or total number of | |
| bytes | |
| value: int (0) | |
| Starting internal counter value | |
| hooks: dict or None | |
| A dict of named functions to be called on each update. The signature | |
| of these must be ``f(size, value, **kwargs)`` | |
| """ | |
| def __init__(self, size=None, value=0, hooks=None, **kwargs): | |
| self.size = size | |
| self.value = value | |
| self.hooks = hooks or {} | |
| self.kw = kwargs | |
| def __enter__(self): | |
| return self | |
| def __exit__(self, *exc_args): | |
| self.close() | |
| def close(self): | |
| """Close callback.""" | |
| def branched(self, path_1, path_2, **kwargs): | |
| """ | |
| Return callback for child transfers | |
| If this callback is operating at a higher level, e.g., put, which may | |
| trigger transfers that can also be monitored. The function returns a callback | |
| that has to be passed to the child method, e.g., put_file, | |
| as `callback=` argument. | |
| The implementation uses `callback.branch` for compatibility. | |
| When implementing callbacks, it is recommended to override this function instead | |
| of `branch` and avoid calling `super().branched(...)`. | |
| Prefer using this function over `branch`. | |
| Parameters | |
| ---------- | |
| path_1: str | |
| Child's source path | |
| path_2: str | |
| Child's destination path | |
| **kwargs: | |
| Arbitrary keyword arguments | |
| Returns | |
| ------- | |
| callback: Callback | |
| A callback instance to be passed to the child method | |
| """ | |
| self.branch(path_1, path_2, kwargs) | |
| # mutate kwargs so that we can force the caller to pass "callback=" explicitly | |
| return kwargs.pop("callback", DEFAULT_CALLBACK) | |
| def branch_coro(self, fn): | |
| """ | |
| Wraps a coroutine, and pass a new child callback to it. | |
| """ | |
| async def func(path1, path2: str, **kwargs): | |
| with self.branched(path1, path2, **kwargs) as child: | |
| return await fn(path1, path2, callback=child, **kwargs) | |
| return func | |
| def set_size(self, size): | |
| """ | |
| Set the internal maximum size attribute | |
| Usually called if not initially set at instantiation. Note that this | |
| triggers a ``call()``. | |
| Parameters | |
| ---------- | |
| size: int | |
| """ | |
| self.size = size | |
| self.call() | |
| def absolute_update(self, value): | |
| """ | |
| Set the internal value state | |
| Triggers ``call()`` | |
| Parameters | |
| ---------- | |
| value: int | |
| """ | |
| self.value = value | |
| self.call() | |
| def relative_update(self, inc=1): | |
| """ | |
| Delta increment the internal counter | |
| Triggers ``call()`` | |
| Parameters | |
| ---------- | |
| inc: int | |
| """ | |
| self.value += inc | |
| self.call() | |
| def call(self, hook_name=None, **kwargs): | |
| """ | |
| Execute hook(s) with current state | |
| Each function is passed the internal size and current value | |
| Parameters | |
| ---------- | |
| hook_name: str or None | |
| If given, execute on this hook | |
| kwargs: passed on to (all) hook(s) | |
| """ | |
| if not self.hooks: | |
| return | |
| kw = self.kw.copy() | |
| kw.update(kwargs) | |
| if hook_name: | |
| if hook_name not in self.hooks: | |
| return | |
| return self.hooks[hook_name](self.size, self.value, **kw) | |
| for hook in self.hooks.values() or []: | |
| hook(self.size, self.value, **kw) | |
| def wrap(self, iterable): | |
| """ | |
| Wrap an iterable to call ``relative_update`` on each iterations | |
| Parameters | |
| ---------- | |
| iterable: Iterable | |
| The iterable that is being wrapped | |
| """ | |
| for item in iterable: | |
| self.relative_update() | |
| yield item | |
| def branch(self, path_1, path_2, kwargs): | |
| """ | |
| Set callbacks for child transfers | |
| If this callback is operating at a higher level, e.g., put, which may | |
| trigger transfers that can also be monitored. The passed kwargs are | |
| to be *mutated* to add ``callback=``, if this class supports branching | |
| to children. | |
| Parameters | |
| ---------- | |
| path_1: str | |
| Child's source path | |
| path_2: str | |
| Child's destination path | |
| kwargs: dict | |
| arguments passed to child method, e.g., put_file. | |
| Returns | |
| ------- | |
| """ | |
| return None | |
| def no_op(self, *_, **__): | |
| pass | |
| def __getattr__(self, item): | |
| """ | |
| If undefined methods are called on this class, nothing happens | |
| """ | |
| return self.no_op | |
| def as_callback(cls, maybe_callback=None): | |
| """Transform callback=... into Callback instance | |
| For the special value of ``None``, return the global instance of | |
| ``NoOpCallback``. This is an alternative to including | |
| ``callback=DEFAULT_CALLBACK`` directly in a method signature. | |
| """ | |
| if maybe_callback is None: | |
| return DEFAULT_CALLBACK | |
| return maybe_callback | |
| class NoOpCallback(Callback): | |
| """ | |
| This implementation of Callback does exactly nothing | |
| """ | |
| def call(self, *args, **kwargs): | |
| return None | |
| class DotPrinterCallback(Callback): | |
| """ | |
| Simple example Callback implementation | |
| Almost identical to Callback with a hook that prints a char; here we | |
| demonstrate how the outer layer may print "#" and the inner layer "." | |
| """ | |
| def __init__(self, chr_to_print="#", **kwargs): | |
| self.chr = chr_to_print | |
| super().__init__(**kwargs) | |
| def branch(self, path_1, path_2, kwargs): | |
| """Mutate kwargs to add new instance with different print char""" | |
| kwargs["callback"] = DotPrinterCallback(".") | |
| def call(self, **kwargs): | |
| """Just outputs a character""" | |
| print(self.chr, end="") | |
| class TqdmCallback(Callback): | |
| """ | |
| A callback to display a progress bar using tqdm | |
| Parameters | |
| ---------- | |
| 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_cls`. | |
| tqdm_cls: (optional) | |
| subclass of `tqdm.tqdm`. If not passed, it will default to `tqdm.tqdm`. | |
| Examples | |
| -------- | |
| >>> import fsspec | |
| >>> from fsspec.callbacks import TqdmCallback | |
| >>> fs = fsspec.filesystem("memory") | |
| >>> path2distant_data = "/your-path" | |
| >>> fs.upload( | |
| ".", | |
| path2distant_data, | |
| recursive=True, | |
| callback=TqdmCallback(), | |
| ) | |
| You can forward args to tqdm using the ``tqdm_kwargs`` parameter. | |
| >>> fs.upload( | |
| ".", | |
| path2distant_data, | |
| 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): | |
| try: | |
| from tqdm import tqdm | |
| except ImportError as exce: | |
| raise ImportError( | |
| "Using TqdmCallback requires tqdm to be installed" | |
| ) from exce | |
| self._tqdm_cls = kwargs.pop("tqdm_cls", tqdm) | |
| self._tqdm_kwargs = tqdm_kwargs or {} | |
| self.tqdm = None | |
| super().__init__(*args, **kwargs) | |
| def call(self, *args, **kwargs): | |
| if self.tqdm is None: | |
| self.tqdm = self._tqdm_cls(total=self.size, **self._tqdm_kwargs) | |
| self.tqdm.total = self.size | |
| self.tqdm.update(self.value - self.tqdm.n) | |
| def close(self): | |
| if self.tqdm is not None: | |
| self.tqdm.close() | |
| self.tqdm = None | |
| def __del__(self): | |
| return self.close() | |
| DEFAULT_CALLBACK = _DEFAULT_CALLBACK = NoOpCallback() | |