import functools
import cProfile
import time
import pstats


def profile(func):
  """Decorator for profiling the execution of a function."""
  @functools.wraps(func)
  def func_wrapped(*args, **kwargs):
    """Function which wraps original function with start/stop profiling."""
    pr = cProfile.Profile()
    pr.enable()
    start = time.time()
    output = func(*args, **kwargs)
    print("Elapsed", time.time() - start)
    pr.disable()
    ps = pstats.Stats(pr).sort_stats("cumulative")
    ps.print_stats()
    return output
  return func_wrapped