| | import contextlib |
| | from typing import Any, Iterable, Iterator, Optional |
| |
|
| | try: |
| | from tqdm import tqdm |
| | except ImportError: |
| | tqdm = None |
| |
|
| | try: |
| | from ray.experimental.tqdm_ray import tqdm as ray_tqdm |
| | except: |
| | ray_tqdm = None |
| |
|
| | |
| | _current_progress_type = "none" |
| | _is_progress_bar_active = False |
| |
|
| |
|
| | class DummyProgressBar: |
| | """A no-op progress bar that mimics tqdm interface""" |
| |
|
| | def __init__(self, iterable=None, **kwargs): |
| | self.iterable = iterable |
| |
|
| | def __iter__(self): |
| | return iter(self.iterable) |
| |
|
| | def update(self, n=1): |
| | pass |
| |
|
| | def close(self): |
| | pass |
| |
|
| | def set_description(self, desc): |
| | pass |
| |
|
| |
|
| | def get_new_progress_bar(iterable: Optional[Iterable] = None, **kwargs) -> Any: |
| | if not _is_progress_bar_active: |
| | return DummyProgressBar(iterable=iterable, **kwargs) |
| |
|
| | if _current_progress_type == "tqdm": |
| | if tqdm is None: |
| | raise ImportError("tqdm is required but not installed. Please install tqdm to use the tqdm progress bar.") |
| | return tqdm(iterable=iterable, **kwargs) |
| | elif _current_progress_type == "ray_tqdm": |
| | if ray_tqdm is None: |
| | raise ImportError("ray is required but not installed. Please install ray to use the ray_tqdm progress bar.") |
| | return ray_tqdm(iterable=iterable, **kwargs) |
| | return DummyProgressBar(iterable=iterable, **kwargs) |
| |
|
| |
|
| | @contextlib.contextmanager |
| | def progress_bar(type: str = "none", enabled=True): |
| | """ |
| | Context manager for setting progress bar type and options. |
| | |
| | Args: |
| | type: Type of progress bar ("none" or "tqdm") |
| | **options: Options to pass to the progress bar (e.g., total, desc) |
| | |
| | Raises: |
| | ValueError: If progress bar type is invalid |
| | RuntimeError: If progress bars are nested |
| | |
| | Example: |
| | with progress_bar(type="tqdm", total=100): |
| | for i in get_new_progress_bar(range(100)): |
| | process(i) |
| | """ |
| | if type not in ("none", "tqdm", "ray_tqdm"): |
| | raise ValueError("Progress bar type must be 'none' or 'tqdm' or 'ray_tqdm'") |
| | if not enabled: |
| | type = "none" |
| | global _current_progress_type, _is_progress_bar_active |
| |
|
| | if _is_progress_bar_active: |
| | raise RuntimeError("Nested progress bars are not supported") |
| |
|
| | _is_progress_bar_active = True |
| | _current_progress_type = type |
| |
|
| | try: |
| | yield |
| | finally: |
| | _is_progress_bar_active = False |
| | _current_progress_type = "none" |
| |
|
