| 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" |
|
|
