codekingpro/portable-devtools

1from functools import wraps2 3 4class Callback:5    """6    Base class and interface for callback mechanism7 8    This class can be used directly for monitoring file transfers by9    providing ``callback=Callback(hooks=...)`` (see the ``hooks`` argument,10    below), or subclassed for more specialised behaviour.11 12    Parameters13    ----------14    size: int (optional)15        Nominal quantity for the value that corresponds to a complete16        transfer, e.g., total number of tiles or total number of17        bytes18    value: int (0)19        Starting internal counter value20    hooks: dict or None21        A dict of named functions to be called on each update. The signature22        of these must be ``f(size, value, **kwargs)``23    """24 25    def __init__(self, size=None, value=0, hooks=None, **kwargs):26        self.size = size27        self.value = value28        self.hooks = hooks or {}29        self.kw = kwargs30 31    def __enter__(self):32        return self33 34    def __exit__(self, *exc_args):35        self.close()36 37    def close(self):38        """Close callback."""39 40    def branched(self, path_1, path_2, **kwargs):41        """42        Return callback for child transfers43 44        If this callback is operating at a higher level, e.g., put, which may45        trigger transfers that can also be monitored. The function returns a callback46        that has to be passed to the child method, e.g., put_file,47        as `callback=` argument.48 49        The implementation uses `callback.branch` for compatibility.50        When implementing callbacks, it is recommended to override this function instead51        of `branch` and avoid calling `super().branched(...)`.52 53        Prefer using this function over `branch`.54 55        Parameters56        ----------57        path_1: str58            Child's source path59        path_2: str60            Child's destination path61        **kwargs:62            Arbitrary keyword arguments63 64        Returns65        -------66        callback: Callback67            A callback instance to be passed to the child method68        """69        self.branch(path_1, path_2, kwargs)70        # mutate kwargs so that we can force the caller to pass "callback=" explicitly71        return kwargs.pop("callback", DEFAULT_CALLBACK)72 73    def branch_coro(self, fn):74        """75        Wraps a coroutine, and pass a new child callback to it.76        """77 78        @wraps(fn)79        async def func(path1, path2: str, **kwargs):80            with self.branched(path1, path2, **kwargs) as child:81                return await fn(path1, path2, callback=child, **kwargs)82 83        return func84 85    def set_size(self, size):86        """87        Set the internal maximum size attribute88 89        Usually called if not initially set at instantiation. Note that this90        triggers a ``call()``.91 92        Parameters93        ----------94        size: int95        """96        self.size = size97        self.call()98 99    def absolute_update(self, value):100        """101        Set the internal value state102 103        Triggers ``call()``104 105        Parameters106        ----------107        value: int108        """109        self.value = value110        self.call()111 112    def relative_update(self, inc=1):113        """114        Delta increment the internal counter115 116        Triggers ``call()``117 118        Parameters119        ----------120        inc: int121        """122        self.value += inc123        self.call()124 125    def call(self, hook_name=None, **kwargs):126        """127        Execute hook(s) with current state128 129        Each function is passed the internal size and current value130 131        Parameters132        ----------133        hook_name: str or None134            If given, execute on this hook135        kwargs: passed on to (all) hook(s)136        """137        if not self.hooks:138            return139        kw = self.kw.copy()140        kw.update(kwargs)141        if hook_name:142            if hook_name not in self.hooks:143                return144            return self.hooks[hook_name](self.size, self.value, **kw)145        for hook in self.hooks.values() or []:146            hook(self.size, self.value, **kw)147 148    def wrap(self, iterable):149        """150        Wrap an iterable to call ``relative_update`` on each iterations151 152        Parameters153        ----------154        iterable: Iterable155            The iterable that is being wrapped156        """157        for item in iterable:158            self.relative_update()159            yield item160 161    def branch(self, path_1, path_2, kwargs):162        """163        Set callbacks for child transfers164 165        If this callback is operating at a higher level, e.g., put, which may166        trigger transfers that can also be monitored. The passed kwargs are167        to be *mutated* to add ``callback=``, if this class supports branching168        to children.169 170        Parameters171        ----------172        path_1: str173            Child's source path174        path_2: str175            Child's destination path176        kwargs: dict177            arguments passed to child method, e.g., put_file.178 179        Returns180        -------181 182        """183        return None184 185    def no_op(self, *_, **__):186        pass187 188    def __getattr__(self, item):189        """190        If undefined methods are called on this class, nothing happens191        """192        return self.no_op193 194    @classmethod195    def as_callback(cls, maybe_callback=None):196        """Transform callback=... into Callback instance197 198        For the special value of ``None``, return the global instance of199        ``NoOpCallback``. This is an alternative to including200        ``callback=DEFAULT_CALLBACK`` directly in a method signature.201        """202        if maybe_callback is None:203            return DEFAULT_CALLBACK204        return maybe_callback205 206 207class NoOpCallback(Callback):208    """209    This implementation of Callback does exactly nothing210    """211 212    def call(self, *args, **kwargs):213        return None214 215 216class DotPrinterCallback(Callback):217    """218    Simple example Callback implementation219 220    Almost identical to Callback with a hook that prints a char; here we221    demonstrate how the outer layer may print "#" and the inner layer "."222    """223 224    def __init__(self, chr_to_print="#", **kwargs):225        self.chr = chr_to_print226        super().__init__(**kwargs)227 228    def branch(self, path_1, path_2, kwargs):229        """Mutate kwargs to add new instance with different print char"""230        kwargs["callback"] = DotPrinterCallback(".")231 232    def call(self, **kwargs):233        """Just outputs a character"""234        print(self.chr, end="")235 236 237class TqdmCallback(Callback):238    """239    A callback to display a progress bar using tqdm240 241    Parameters242    ----------243    tqdm_kwargs : dict, (optional)244        Any argument accepted by the tqdm constructor.245        See the `tqdm doc <https://tqdm.github.io/docs/tqdm/#__init__>`_.246        Will be forwarded to `tqdm_cls`.247    tqdm_cls: (optional)248        subclass of `tqdm.tqdm`. If not passed, it will default to `tqdm.tqdm`.249 250    Examples251    --------252    >>> import fsspec253    >>> from fsspec.callbacks import TqdmCallback254    >>> fs = fsspec.filesystem("memory")255    >>> path2distant_data = "/your-path"256    >>> fs.upload(257            ".",258            path2distant_data,259            recursive=True,260            callback=TqdmCallback(),261        )262 263    You can forward args to tqdm using the ``tqdm_kwargs`` parameter.264 265    >>> fs.upload(266            ".",267            path2distant_data,268            recursive=True,269            callback=TqdmCallback(tqdm_kwargs={"desc": "Your tqdm description"}),270        )271 272    You can also customize the progress bar by passing a subclass of `tqdm`.273 274    .. code-block:: python275 276        class TqdmFormat(tqdm):277            '''Provides a `total_time` format parameter'''278            @property279            def format_dict(self):280                d = super().format_dict281                total_time = d["elapsed"] * (d["total"] or 0) / max(d["n"], 1)282                d.update(total_time=self.format_interval(total_time) + " in total")283                return d284 285    >>> with TqdmCallback(286            tqdm_kwargs={287                "desc": "desc",288                "bar_format": "{total_time}: {percentage:.0f}%|{bar}{r_bar}",289            },290            tqdm_cls=TqdmFormat,291        ) as callback:292            fs.upload(".", path2distant_data, recursive=True, callback=callback)293    """294 295    def __init__(self, tqdm_kwargs=None, *args, **kwargs):296        try:297            from tqdm import tqdm298 299        except ImportError as exce:300            raise ImportError(301                "Using TqdmCallback requires tqdm to be installed"302            ) from exce303 304        self._tqdm_cls = kwargs.pop("tqdm_cls", tqdm)305        self._tqdm_kwargs = tqdm_kwargs or {}306        self.tqdm = None307        super().__init__(*args, **kwargs)308 309    def call(self, *args, **kwargs):310        if self.tqdm is None:311            self.tqdm = self._tqdm_cls(total=self.size, **self._tqdm_kwargs)312        self.tqdm.total = self.size313        self.tqdm.update(self.value - self.tqdm.n)314 315    def close(self):316        if self.tqdm is not None:317            self.tqdm.close()318            self.tqdm = None319 320    def __del__(self):321        return self.close()322 323 324DEFAULT_CALLBACK = _DEFAULT_CALLBACK = NoOpCallback()325