import time


def timeit(f):
    """
    Modify a function to print cpu and wall-clock elapsed time when called.

    Can be used as a decorator:
    @timeit
    def func(x):
        ...
    """
    def g(*args, **kwargs):
        start_etime = time.perf_counter()
        start_cputime = time.process_time()
        rvalue = f(*args, **kwargs)
        end_etime = time.perf_counter()
        end_cputime = time.process_time()
        print('elapsed time (s): ', end_etime - start_etime)
        print('cpu time (s)', end_cputime - start_cputime)
        return rvalue
    return g