Source code for archai.common.stopwatch

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

# adapted from https://github.com/ildoonet/pystopwatch2/blob/master/pystopwatch2/watch.py

import threading
import time
from collections import defaultdict
from enum import Enum


class _ClockState(Enum):
    PAUSE = 0
    RUN = 1

class _Clock:
    tag_default = '__default1958__'
    th_lock = threading.Lock()

    def __init__(self):
        self.prev_time = time.time()
        self.sum = 0.
        self.state = _ClockState.PAUSE

    def __str__(self):
        return 'state=%s elapsed=%.4f prev_time=%.8f' % (self.state, self.sum, self.prev_time)

    def __repr__(self):
        return self.__str__()


class StopWatch:
    stopwatch:'StopWatch' = None

    def __init__(self):
        self.clocks = defaultdict(lambda: _Clock())

    def start(self, tag=None):
        if tag is None:
            tag = _Clock.tag_default
        with _Clock.th_lock:
            clock = self.clocks[tag]
            if clock.state == _ClockState.RUN:
                return
            clock.state = _ClockState.RUN
            clock.prev_time = time.time()

    def pause(self, tag=None):
        if tag is None:
            tag = _Clock.tag_default
        with _Clock.th_lock:
            clock = self.clocks[tag]
            clock.state = _ClockState.PAUSE
            delta = time.time() - clock.prev_time
            clock.sum += delta
            return clock.sum

    def clear(self, tag=None):
        if tag is None:
            tag = _Clock.tag_default
        del self.clocks[tag]

    def get_elapsed(self, tag=None):
        if tag is None:
            tag = _Clock.tag_default
        clock = self.clocks[tag]
        elapsed = clock.sum
        if clock.state == _ClockState.RUN:
            elapsed += time.time() - clock.prev_time

        return elapsed

    def keys(self):
        return self.clocks.keys()

    def __str__(self):
        return '\n'.join(['%s: %s' % (k, v) for k, v in self.clocks.items()])

    def __repr__(self):
        return self.__str__()

    @staticmethod
    def set(instance:'StopWatch')->None:
        StopWatch.stopwatch = instance

    @staticmethod
    def get()->'StopWatch':
        return StopWatch.stopwatch