A benchmarking decorator for easy use of timeit

enable_timeit.py

def enable_timeit(number=1000, repeat=1):
    """
    Adds a `timeit` member function to the decorated function.

    This function allows for configuration of how `timeit` will be executed for
    the decorated function. Note that ultimately, `timeit.repeat()` will be
    called, not `timeit.timeit()`.

    Args:
        number (int): How often `timeit` should execute the function.
        repeat (int): How often to repeat the entire `timeit` loop.

    Returns:
        The actual decorator.
    """

    def inner_enable_timeit(function):
        """Adds a `timeit` function to the decorated function."""

        # The actual wrapped function (identity function)
        def proxy(*args, **kwargs):
            return function(*args, **kwargs)

        def execute_timeit(*args, **kwargs):
            """
            Times the function with the given arguments.

            Args:
                args (list): The positional arguments to forward.
                kwargs (dict): The keyword arguments to forward.
            """

            # Stringify the positional and keyword arguments
            args = ', '.join(map(str, args))
            kwargs = ', '.join('{0}={1}'.format(k,v) for k,v in kwargs.items())
            if args and kwargs:
                arguments = ', '.join((args, kwargs))
            else:
                # At least one will be the empty string
                arguments = args + kwargs
            statement = '{0}({1})'.format(function.__name__, arguments)

            result = timeit.repeat(
                statement,
                number=number,
                repeat=repeat,
                globals={function.__name__: function}
            )

            # Unpack the single result for convenience
            if repeat == 1:
                return result[0]
            return result

        # Add the "member function"
        proxy.timeit = execute_timeit

        return proxy

    return inner_enable_timeit