from typing import Callable, Iterable, Union, Any from tqdm import tqdm def njit(f: Union[Callable, str] = None, *args, **kwargs) -> Callable: """Wrapper for optional numba's njit decorator Args: f (Union[Callable, str], optional): Function to wrap with numba. Defaults to None. Returns: Callable: Wrapped function. """ def decorator(func: Callable) -> Any: return func if callable(f): return f return decorator def tqdm_iter(iter: Iterable, desc: str): """Wrapper for optional tqdm iterator progress bar. Args: iter (Iterable): Object to iterate over. desc (str): Description written to stdout. Returns: _type_: Wrapped iterator. """ return tqdm(iter, leave = False, desc = desc)