30 lines
713 B
Python
30 lines
713 B
Python
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) |