Source code for aiida.common.progress_reporter

###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""Provide a singleton progress reporter implementation.

The interface is inspired by `tqdm <https://github.com/tqdm/tqdm>`,
and indeed a valid implementation is::

    from tqdm import tqdm
    set_progress_reporter(tqdm, bar_format='{l_bar}{bar}{r_bar}')

"""
from functools import partial
from types import TracebackType
from typing import Any, Callable, Optional, Type

__all__ = (
    'get_progress_reporter',
    'set_progress_reporter',
    'set_progress_bar_tqdm',
    'ProgressReporterAbstract',
    'TQDM_BAR_FORMAT',
    'create_callback',
)

TQDM_BAR_FORMAT = '{desc:40.40}{percentage:6.1f}%|{bar}| {n_fmt}/{total_fmt}'


[docs] class ProgressReporterAbstract: """An abstract class for incrementing a progress reporter. This class provides the base interface for any `ProgressReporter` class. Example Usage:: with ProgressReporter(total=10, desc="A process:") as progress: for i in range(10): progress.set_description_str(f"A process: {i}") progress.update() """
[docs] def __init__(self, *, total: int, desc: Optional[str] = None, **kwargs: Any): """Initialise the progress reporting contextmanager. :param total: The number of expected iterations. :param desc: A description of the process """ self._total = total self._desc = desc self._increment: int = 0
@property def total(self) -> int: """Return the total iterations expected.""" return self._total @property def desc(self) -> Optional[str]: """Return the description of the process.""" return self._desc @property def n(self) -> int: """Return the current iteration.""" # note using `n` as the attribute name is necessary for compatibility with tqdm return self._increment
[docs] def __enter__(self) -> 'ProgressReporterAbstract': """Enter the contextmanager.""" return self
[docs] def __exit__( self, exctype: Optional[Type[BaseException]], excinst: Optional[BaseException], exctb: Optional[TracebackType] ): """Exit the contextmanager.""" return False
[docs] def set_description_str(self, text: Optional[str] = None, refresh: bool = True): """Set the text shown by the progress reporter. :param text: The text to show :param refresh: Force refresh of the progress reporter """ self._desc = text
[docs] def update(self, n: int = 1): """Update the progress counter. :param n: Increment to add to the internal counter of iterations """ self._increment += n
[docs] def reset(self, total: Optional[int] = None): """Resets current iterations to 0. :param total: If not None, update number of expected iterations. """ self._increment = 0 if total is not None: self._total = total
[docs] class ProgressReporterNull(ProgressReporterAbstract): """A null implementation of the progress reporter. This implementation does not output anything. """
PROGRESS_REPORTER: Type[ProgressReporterAbstract] = ProgressReporterNull
[docs] def get_progress_reporter() -> Type[ProgressReporterAbstract]: """Return the progress reporter Example Usage:: with get_progress_reporter()(total=10, desc="A process:") as progress: for i in range(10): progress.set_description_str(f"A process: {i}") progress.update() """ global PROGRESS_REPORTER # noqa: PLW0602 return PROGRESS_REPORTER
[docs] def set_progress_reporter(reporter: Optional[Type[ProgressReporterAbstract]] = None, **kwargs: Any): """Set the progress reporter implementation :param reporter: A progress reporter for a process. If None, reset to ``ProgressReporterNull``. :param kwargs: If present, set a partial function with these kwargs The reporter should be a context manager that implements the :func:`~aiida.common.progress_reporter.ProgressReporterAbstract` interface. Example Usage:: set_progress_reporter(ProgressReporterNull) with get_progress_reporter()(total=10, desc="A process:") as progress: for i in range(10): progress.set_description_str(f"A process: {i}") progress.update() """ global PROGRESS_REPORTER # noqa: PLW0603 if reporter is None: PROGRESS_REPORTER = ProgressReporterNull elif kwargs: PROGRESS_REPORTER = partial(reporter, **kwargs) # type: ignore[assignment] else: PROGRESS_REPORTER = reporter
[docs] def set_progress_bar_tqdm(bar_format: Optional[str] = TQDM_BAR_FORMAT, leave: Optional[bool] = False, **kwargs: Any): """Set a `tqdm <https://github.com/tqdm/tqdm>`__ implementation of the progress reporter interface. See :func:`~aiida.common.progress_reporter.set_progress_reporter` for details. :param bar_format: Specify a custom bar string format. :param leave: If True, keeps all traces of the progressbar upon termination of iteration. If `None`, will leave only if `position` is `0`. :param kwargs: pass to the tqdm init """ from tqdm import tqdm set_progress_reporter(tqdm, bar_format=bar_format, leave=leave, **kwargs)
[docs] def create_callback(progress_reporter: ProgressReporterAbstract) -> Callable[[str, Any], None]: """Create a callback function to update the progress reporter. :returns: a callback to report on the process, ``callback(action, value)``, with the following callback signatures: - ``callback('init', {'total': <int>, 'description': <str>})``, to reset the progress with a new total iterations and description - ``callback('update', <int>)``, to update the progress by a certain number of iterations """ def _callback(action: str, value: Any): if action == 'init': progress_reporter.reset(value['total']) progress_reporter.set_description_str(value['description']) elif action == 'update': progress_reporter.update(value) return _callback