Source code for schrodinger.application.desmond.packages.timer

"""
Facilities for timing a block of code.
Copyright Schrodinger, LLC. All rights reserved.

Example usage:

   from schrodinger.application.desmond.packages.timer import timer

   with timer("Reading the trajectory..."):
       traj.read_trj(trj_fname)

Example output:

   Reading the trajectory...
     100 frames
     Done. (0.5 sec)

In this example, we have three lines in output. Line#1 is the start message, and
line#3 is the timing message. The start-message and the format for the timing
message can be customed by the `message` and `fmt` arguments, respectively.

Arguments for `timer`:

:type    message: `str` or `None`
:type    message: The message to print out before the timing starts. Default is
                  `None` -- no message will be printed.
:type        fmt: `str`
:type        fmt: The format to print out the timing. Default is
                  "Done. (%.2f sec)". A custom format should contain 1 floating
                  number control paramter.
:type  threshold: `float`
:param threshold: Specifies a threshold. If the elapsed time is less than
                  `threshold`, no timing message will not be printed out.
:type     record: `list` or `None`
:param    record: If a list value is given, the current elapsed time will be
                  appended to the list. This is handy if you want to accumulate
                  the elapsed times.
:type     logger: A callable object.
:param    logger: Function called to print the messages.
"""

import time
from collections import namedtuple
from contextlib import contextmanager


def _default_logger(message):
    print(message)


class Timer(object):
    """
    You can use this class to create your own timer function with custom default
    behavior.

    Example usage:

      from schrodinger.application.desmond.packages.timer import Timer

      mytimer = Timer(threshold=0.2)
      # Creates a timer with 0.2 threshold.

      with mytimer("Reading the trajectory..."):
        traj.read_trj(trj_fname)

    In this example, `mytimer` behaves exactly the same as the standard
    `timer` (see the docstring of this module) except for the default value of
    `threshold` parameter that has been customed to 0.2.
    """

    Args = namedtuple("Args", "message threshold fmt record logger")

    def __init__(self,
                 message=None,
                 threshold=0.1,
                 fmt="  Done. (%.2f sec)",
                 record=None,
                 logger=_default_logger):
        self._default = Timer.Args(message, threshold, fmt, record, logger)

    @contextmanager
    def __call__(self, *arg, **kwarg):
        a = list(self._default)
        a = list(arg) + a[len(arg):]  # Updates `a' with `arg'.
        a = Timer.Args(*a)
        a = a._replace(**kwarg)  # Updates `a' with `kwarg'.
        message, threshold, fmt, record, logger = a
        message and logger(message)
        start_time = time.time()
        yield
        elapsed_time = time.time() - start_time
        if elapsed_time > threshold:
            logger(fmt % elapsed_time)
        if isinstance(record, list):
            record.append(elapsed_time)


timer = Timer()