crossbench/helper/durations.py

# Copyright 2024 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from __future__ import annotations

import contextlib
import datetime as dt
import logging
from typing import Iterator, Self


class Durations:
  """
  Helper object to track durations.
  """

  def __init__(self) -> None:
    self._durations: dict[str, dt.timedelta] = {}

  def __getitem__(self, name: str) -> dt.timedelta:
    return self._durations[name]

  def __setitem__(self, name: str, duration: dt.timedelta) -> None:
    assert name not in self._durations, f"Cannot set '{name}' duration twice!"
    self._durations[name] = duration

  def __len__(self) -> int:
    return len(self._durations)

  @contextlib.contextmanager
  def measure(self, name: str) -> Iterator[None]:
    assert name not in self._durations, (
        f"Cannot measure '{name}' duration twice!")
    start_time = dt.datetime.now()
    try:
      yield
    finally:
      delta = dt.datetime.now() - start_time
      self._durations[name] = delta

  def to_json(self) -> dict[str, float]:
    return {
        name: self._durations[name].total_seconds()
        for name in sorted(self._durations.keys())
    }


class TimeScope:
  """
  Measures and logs the time spend during the lifetime of the TimeScope.
  """

  def __init__(self, message: str, level: int = 3) -> None:
    self._message = message
    self._level = level
    self._start: dt.datetime | None = None
    self._duration: dt.timedelta = dt.timedelta()

  @property
  def message(self) -> str:
    return self._message

  @property
  def duration(self) -> dt.timedelta:
    return self._duration

  def __enter__(self) -> Self:
    self._start = dt.datetime.now()
    return self

  def __exit__(self, exc_type, exc_value, traceback) -> None:  # noqa: ANN001
    assert self._start
    self._duration = dt.datetime.now() - self._start
    logging.log(self._level, "%s duration=%s", self._message, self._duration)