timer.py

import time
import json


class Timer:
    def __init__(self, name=None):
        self.name = name
        self.start_time = None
        self.steps = []
        self.total_time = None

    def clear(self):
        self.start_time = None
        self.steps = []
        self.total_time = None

    def start(self):
        """Start the timer."""
        self.start_time = time.time()

    def is_running(self):
        return self.start_time is not None

    def add_step(self, step_name):
        """Add a step with its duration since the last step or start."""
        if self.start_time is None:
            self.start()

        current_time = time.time()
        if not self.steps:
            elapsed = current_time - self.start_time
        else:
            elapsed = current_time - self.steps[-1]['timestamp']

        self.steps.append({
            "step_name": step_name,
            "duration": round(elapsed, 4),
            "total_duration": round(current_time - self.start_time, 4),
            "timestamp": current_time
        })

    def end(self):
        """End the timer and calculate the total duration."""
        if self.start_time is None:
            raise RuntimeError("Timer has not been started.")

        if not self.steps:
            raise RuntimeError("No steps have been added.")

        self.total_time = time.time() - self.start_time

    def to_json(self):
        """Return a JSON of the timing steps."""
        if self.total_time is None:
            raise RuntimeError("Timer has not been ended.")

        output_steps = {}
        for step in self.steps:
            output_steps[step["step_name"]] = step["duration"]

        highlights = {"total_time": round(self.total_time, 4)}

        if self.name:
            highlights = {"name": self.name, **highlights}

        output = {
            **highlights,
            **output_steps
        }
        return output

    def to_json_str(self):
        """Return a human-readable JSON of the timing steps."""
        return json.dumps(self.to_json(), indent=4)

    def formatted_result(self):
        """Return a list of the steps, their duration, and total duration."""
        if self.total_time is None:
            raise RuntimeError("Timer has not been ended.")
        line_buffer = []
        if self.name:
            line_buffer.append(f"Timer: {self.name}")
        for step in self.steps:
            line_buffer.append(f"[{step['duration']:05.2f}s, {step['total_duration']:05.2f}s] {step['step_name']}")
        # for step in self.steps:
        #     line_buffer.append(f"{step['step_name']}: {step['duration']:.2f}s ({step['total_duration']:.2f}s)")
        line_buffer.append(f"Total time: {self.total_time:.2f}s")
        return "\n".join(line_buffer)

    def log_formatted_result(self):
        print(self.formatted_result())


def example():
    # Example usage
    timer = Timer()
    timer.start()

    # Simulating some steps
    time.sleep(1)  # Simulate work for step 1
    timer.add_step("Step 1")

    time.sleep(2)  # Simulate work for step 2
    timer.add_step("Step 2")

    timer.end()

    # Print the timer output
    print(timer.formatted_result())
    print(timer.to_json_str())


if __name__ == "__main__":
    example()