Source code for schrodinger.utils.timeit

from functools import wraps
import logging
from time import perf_counter, strftime


class TimeIt:
    """
    Track and report wall time elapsed during a function or code block
    execution. Can be used as a function decorator or context manager.

    Both the context manager and decorator will output time-related messages to
    the log at the specified logging level. The context manager can be assigned
    to a variable, which can be accessed inside the context to get elapsed time
    and log checkpoints. It can also be inspected after the context is closed
    to get total time elapsed.

    Before we get into examples, let's set up our imports and logger
    boilerplate::

        >>> from schrodinger.utils import log
        >>> from schrodinger.utils.timeit import TimeIt
        >>>
        >>> logger = log.get_output_logger(__name__)
        >>> logger.setLevel(log.DEBUG)

    Now, to use this class as a context manager::

        >>> with TimeIt(activity='do foo', logger=logger, level=log.DEBUG) as t:
        >>>     for _ in range(100):
        >>>         # Do something expensive here
        >>>         pass
        >>>     # Write a checkpoint message to log
        >>>     t.checkpoint('first thing done')
        >>>     for _ in range(100):
        >>>         # Do something else here
        >>>         pass
        TimeIt - do foo - Starting - 2022-05-25 01:45:35
        TimeIt - do foo - Checkpoint - first thing done - 0.0000 s
        TimeIt - do foo - Completed - 0.0002 s

    Elapsed time does not update after context is closed::

        >>> assert t.elapsed == t.end_time - t.start_time

    Alternatively, to use the class as a decorator::

        >>> @TimeIt(logger=logger, level=log.INFO)
        >>> def do_foo():
        >>>     for _ in range(100):
        >>>         pass
        >>> do_foo()
        TimeIt - do_foo - Starting - 2022-05-25 01:47:50
        TimeIt - do_foo - Completed - 0.0000 s

    :cvar TIME_FORMAT: The time format string used by `strftime` when reporting
            the time that measurement began.
    :ivar start_time: The time in fractional seconds just prior to entering the
            context. The reference point is undefined. It is only used to
            compare with `end_time`.
    :ivar end_time: The time in fractional seconds just after exiting the
            context. The reference is undefined.
    """
    TIME_FORMAT: str = '%Y-%m-%d %H:%M:%S'
    start_time: float
    end_time: float
    _elapsed: float

    def __init__(self, *, activity=None, logger=None, level=logging.INFO):
        """
        :param str activity: A description of the activity to be timed. If this
                class is used as a decorator and `activity` is NoneType, the
                function name will be used.
        :param logging.Logger logger: A logger to receive timing messages.
        :param int level: The logger level to use for timing messages.
        """
        self.activity = activity
        self.logger = logger
        self.level = level
        self._checkpoint = 0
        self._has_exited = False

    def __call__(self, func):
        if self.activity is None:
            self.activity = func.__name__

        # Just like contextlib.ContextDecorator
        @wraps(func)
        def inner(*args, **kwds):
            with self:
                return func(*args, **kwds)

        return inner

    @property
    def elapsed(self) -> float:
        """
        Fractional seconds elapsed since context was entered. If context manager
        has already exited, this is the total execution time of the code block
        inside the context manager.

        :return: Fractional seconds elapsed since the start time.
        """
        return self._elapsed if self._has_exited \
            else perf_counter() - self.start_time

    def checkpoint(self, description=None) -> float:
        """
        Write a checkpoint to the log and get elapsed time. Be aware that
        logging can add non-negligible time to a code block's execution.
        Empirically, each call to this function appears add about 0.5ms.

        :param description: A description of the checkpoint. If not supplied,
            the sequential checkpoint number will be used.
        :return: Fractional seconds elapsed since the start time.
        """
        elapsed = self.elapsed
        if description is None:
            description = self._checkpoint
        self._checkpoint += 1
        self._log('Checkpoint', description, f'{elapsed:.4f} s')
        return elapsed

    def _generate_log_entry(self, items):
        cls_name = self.__class__.__name__
        activity = self.activity or ''
        return ' - '.join(map(str, [cls_name, activity, *items]))

    def _log(self, *items):
        if self.logger:
            entry = self._generate_log_entry(items)
            self.logger.log(self.level, entry)

    def __enter__(self):
        self._log('Starting', strftime(self.TIME_FORMAT))
        self.start_time = perf_counter()
        return self

    def __exit__(self, *exc):
        self.end_time = perf_counter()
        self._elapsed = self.end_time - self.start_time
        self._log('Completed', f'{self._elapsed:.4f} s')
        self._has_exited = True
        return False