Source code for schrodinger.utils.profiling

"""
Profiling and timing decorators.

For more information on profiling python code, see also
https://internal.schrodinger.com/products/development/docs/software_engineering/profiling.html#profiling-python-code

A very simple example showing the use the profiling and timing decorators::

    import time

    @profiling
    def profiled_function(duration=0.05):
        time.sleep(duration)


    @timing
    def timed_function(duration=0.05):
        time.sleep(duration)


    class Test(object):
        @profiling
        def profiledMethod(self, duration=0.05):
            time.sleep(duration)

        @timing
        def timedMethod(self, duration=0.05):
            time.sleep(duration)


    test = Test()
    # call the profiled method
    test.profiledMethod()

    # call the timed method
    test.timedMethod(0.1)

    # call the profiled function
    profiled_function()

    # call the timed function
    timed_function(0.1)
"""

import atexit
import contextlib
import cProfile
import datetime
import time
import os

import decorator


@decorator.decorator
def profiling(f, *args, **kwargs):
    """
    Decorator that will generate a profile file with the name of the decorated
    function or method and time stamp.  The duration of the execution of
    decorated function or method will also be printed to the terminal.
    """
    now = datetime.datetime.now()
    filename = now.strftime(f.__name__ + "_%Y%m%d-%H%M%S.%f.prof")
    print("%s profile" % os.path.abspath(filename))
    retval = profile_call(f, *args, **kwargs, profile_filename=filename)
    print("{}: {} sec".format(f.__name__, datetime.datetime.now() - now))
    return retval


def timing(f):
    """
    Decorator that will print the duration of the execution of decorated
    function or method to the terminal every time it's called. Will also
    report the average time, minimum time, and the total number of calls on
    program exit.
    """
    f._timings = []

    def wrapped_f(*args, **kwargs):
        start = time.time()
        retval = f(*args, **kwargs)
        elapsed_time = time.time() - start
        print(f"{f.__name__}: {elapsed_time:.3f} sec")
        f._timings.append(elapsed_time)
        return retval

    atexit.register(_report_average_timing, f)
    return wrapped_f


def _report_average_timing(f):
    if f._timings:
        avg_time = sum(f._timings) / len(f._timings)
        min_time = min(f._timings)
        max_time = max(f._timings)
    else:
        avg_time = 0
        min_time = 0
        max_time = 0

    print(f"TIME [MIN/AVG/MAX] for {f.__name__} over {len(f._timings)} calls:"
          f" {min_time:.3f}/{avg_time:.3f}/{max_time:.3f} sec")


def profile_call(func, *args, profile_filename=None, **kwargs):
    """
    Profile the specified function call.

    :param func: The function to profile.  All arguments to `profile_call` other
        than `profile_filename` will be passed to this function.
    :type func: Callable

    :param profile_filename: The name of the file to save the profiling data to.
        Note that this argument is keyword-only and is required.
    :type profile_filename: str

    :return: The value returned by `func`.
    :rtype: object
    """
    if profile_filename is None:
        raise ValueError("No filename given.")
    profiler = cProfile.Profile()
    retval = profiler.runcall(func, *args, **kwargs)
    profiler.dump_stats(profile_filename)
    return retval


@contextlib.contextmanager
def profile(profile_filename='program.prof'):
    """
    Profile a block of code and dump the results into a file. Should be used as
    a context manager, e.g.,

    with profile(profile_filename='mystats.prof'):
        execute_some_code()

    :param str profile_filename: The name of the file to save the profiling
        data to.

    :rtype: cProfile.Profile
    :return: The profiling object being used
    """

    profiler = cProfile.Profile()
    profiler.enable()
    yield profiler

    profiler.disable()
    profiler.dump_stats(profile_filename)