Python decorator for printing function progresses

print_progress.py

import functools
import inspect
from tqdm import tqdm


def print_progress(total=None, enabled=True):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            if not enabled or total is None:
                return func(*args, **kwargs)
            
            argspec = inspect.getfullargspec(func)
            if 'step_callback' not in argspec.kwonlyargs:
                raise TypeError('{} must have a keyword-only arugment named "step_callback"'.format(func.__name__))
            
            with tqdm(total=total) as pbar:
                step_callback = lambda steps=1: pbar.update(steps)
                ret = func(*args, **kwargs, step_callback=step_callback)
                return ret
        
        return wrapper
    return decorator



# Example
@print_progress(total=10)
def foo(*, step_callback):
    for i in range(10):
        step_callback()