Spaces:
Running
Running
File size: 9,210 Bytes
b72ab63 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 |
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.
"""
@wraps(fn)
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
@classmethod
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()
|