"""
Progress
--------
This used ``tqdm`` to provide a progress bar while computing the metacells.
"""
import logging
from contextlib import contextmanager
from typing import Any
from typing import Dict
from typing import Optional
from typing import Tuple
from tqdm import tqdm # type: ignore
import metacells.utilities.logging as utl
__all__ = [
"progress_bar",
"progress_bar_slice",
"did_progress",
"has_progress_bar",
"start_progress_bar",
"end_progress_bar",
"start_progress_bar_slice",
"end_progress_bar_slice",
]
# Parameters for the progress bar.
TQDM_KWARGS: Optional[Dict[str, Any]] = None
# Global progress bar object.
PROGRESS_BAR: Optional[tqdm] = None
# The log level when the progress bar was activated.
PROGRESS_LOG_LEVEL = logging.INFO
# The current progress position.
PROGRESS_POSITION = 0
# The current progress bar (slice) final position.
PROGRESS_END = 0
#: The size of the current progress bar (slice)
PROGRESS_SIZE = 0
[docs]
def has_progress_bar() -> bool:
"""
Return whether there is an active progress bar.
"""
return PROGRESS_BAR is not None or TQDM_KWARGS is not None
[docs]
@contextmanager
def progress_bar(**tqdm_kwargs: Any) -> Any:
"""
Run some code with a ``tqdm`` progress bar.
..note::
When a progress bar is active, logging is restricted to warnings and errors.
"""
start_progress_bar(**tqdm_kwargs)
try:
result = yield None
return result
finally:
end_progress_bar()
[docs]
def start_progress_bar(**tqdm_kwargs: Any) -> Any:
"""
Create a progress bar (but do not show it yet).
"""
global TQDM_KWARGS
global PROGRESS_POSITION
global PROGRESS_END
global PROGRESS_SIZE
global PROGRESS_LOG_LEVEL
assert PROGRESS_BAR is None
assert TQDM_KWARGS is None
TQDM_KWARGS = tqdm_kwargs.copy()
if "bar_format" not in TQDM_KWARGS:
TQDM_KWARGS["bar_format"] = "{l_bar}{bar}[{elapsed}]" # NOT F-STRING
PROGRESS_POSITION = 0
PROGRESS_END = int(1e15)
PROGRESS_SIZE = PROGRESS_END
PROGRESS_LOG_LEVEL = utl.logger().level
utl.logger().setLevel(logging.WARNING)
[docs]
def end_progress_bar() -> None:
"""
End an active progress bar.
"""
global PROGRESS_BAR
if PROGRESS_BAR is not None:
PROGRESS_BAR.close()
utl.logger().setLevel(PROGRESS_LOG_LEVEL)
PROGRESS_BAR = None
[docs]
@contextmanager
def progress_bar_slice(fraction: Optional[float]) -> Any:
"""
Run some code which will use a slice of the current progress bar.
This can be nested to split the overall progress bar into smaller and smaller parts to represent a tree of
computations.
If ``fraction`` is None, or there is no active progress bar, simply runs the code.
"""
_show_progress_bar()
if fraction is None:
result = yield None
return result
old_state = start_progress_bar_slice(fraction)
try:
result = yield None
return result
finally:
end_progress_bar_slice(old_state)
[docs]
def start_progress_bar_slice(fraction: float) -> Tuple[int, int]:
"""
Start a nested slice of the overall progress bar.
Returned the captured state that needs to be passed to ``end_progress_bar_slice``.
"""
assert PROGRESS_BAR is not None
assert TQDM_KWARGS is None
assert 0.0 < fraction <= 1.0, f"invalid progress fraction: {fraction} (1 + {fraction - 1})"
global PROGRESS_END
global PROGRESS_SIZE
old_progress_size = PROGRESS_SIZE
old_progress_end = PROGRESS_END
PROGRESS_SIZE = int(PROGRESS_SIZE * fraction)
PROGRESS_END = PROGRESS_POSITION + PROGRESS_SIZE
assert (
0 < PROGRESS_SIZE <= old_progress_size
), f"invalid progress size: {PROGRESS_SIZE} ({old_progress_size} + {PROGRESS_SIZE - old_progress_size})"
assert (
PROGRESS_END <= old_progress_end
), f"invalid progress end: {PROGRESS_END} ({old_progress_end} + {PROGRESS_END - old_progress_end})"
return old_progress_size, old_progress_end
[docs]
def end_progress_bar_slice(old_state: Tuple[int, int]) -> None:
"""
End a nested slice of the overall progress bar.
This moves the progress bar position to the end of the slice regardless of reported progress within it.
"""
old_progress_size, old_progress_end = old_state
global PROGRESS_POSITION
global PROGRESS_END
global PROGRESS_SIZE
assert PROGRESS_BAR is not None
assert TQDM_KWARGS is None
assert PROGRESS_POSITION <= PROGRESS_END
remaining = PROGRESS_END - PROGRESS_POSITION
if remaining > 0:
PROGRESS_BAR.update(remaining)
PROGRESS_POSITION = PROGRESS_END
PROGRESS_SIZE = old_progress_size
PROGRESS_END = old_progress_end
[docs]
def did_progress(fraction: float) -> Any:
"""
Report progress of some fraction of the current (slice of) progress bar.
"""
_show_progress_bar()
assert 0.0 < fraction <= 1.0, f"invalid progress fraction: {fraction} (1 + {fraction - 1})"
assert PROGRESS_BAR is not None
assert TQDM_KWARGS is None
global PROGRESS_POSITION
step = int(PROGRESS_SIZE * fraction)
step = min(step, PROGRESS_END - PROGRESS_POSITION)
if step > 0:
PROGRESS_BAR.update(step)
PROGRESS_POSITION += step
def _show_progress_bar() -> None:
global PROGRESS_BAR
global TQDM_KWARGS
if PROGRESS_BAR is None and TQDM_KWARGS is not None:
PROGRESS_BAR = tqdm(total=PROGRESS_SIZE, **TQDM_KWARGS) # pylint: disable=not-a-mapping
TQDM_KWARGS = None