google_nest_sdm.diagnostics

  1"""Diagnostics for debugging."""
  2
  3from __future__ import annotations
  4
  5import time
  6from collections import Counter
  7from collections.abc import Mapping
  8from contextlib import contextmanager
  9from typing import Any, Generator, TypeVar, cast
 10
 11__all__ = [
 12    "get_diagnostics",
 13]
 14
 15
 16class Diagnostics:
 17    """Information for the library."""
 18
 19    def __init__(self) -> None:
 20        """Initialize Diagnostics."""
 21        self._counter: Counter = Counter()
 22        self._subkeys: dict[str, Diagnostics] = {}
 23
 24    def increment(self, key: str, count: int = 1) -> None:
 25        """Increment a counter for the specified key/event."""
 26        self._counter.update(Counter({key: count}))
 27
 28    def elapsed(self, key_prefix: str, elapsed_ms: int = 1) -> None:
 29        """Track a latency event for the specified key/event prefix."""
 30        self.increment(f"{key_prefix}_count", 1)
 31        self.increment(f"{key_prefix}_sum", elapsed_ms)
 32
 33    def as_dict(self) -> Mapping[str, Any]:
 34        """Return diagnostics as a debug dictionary."""
 35        data: dict[str, Any] = {k: self._counter[k] for k in self._counter}
 36        for k, d in self._subkeys.items():
 37            v = d.as_dict()
 38            if not v:
 39                continue
 40            data[k] = v
 41        return data
 42
 43    def subkey(self, key: str) -> Diagnostics:
 44        """Return sub-Diagnositics object with the specified subkey."""
 45        if key not in self._subkeys:
 46            self._subkeys[key] = Diagnostics()
 47        return self._subkeys[key]
 48
 49    @contextmanager
 50    def timer(self, key_prefix: str) -> Generator[None, None, None]:
 51        """A context manager that records the timing of operations as a diagnostic."""
 52        start = time.perf_counter()
 53        try:
 54            yield
 55        finally:
 56            end = time.perf_counter()
 57            ms = int((end - start) * 1000)
 58            self.elapsed(key_prefix, ms)
 59
 60    def reset(self) -> None:
 61        """Clear all diagnostics, for testing."""
 62        self._counter = Counter()
 63        for d in self._subkeys.values():
 64            d.reset()
 65
 66
 67SUBSCRIBER_DIAGNOSTICS = Diagnostics()
 68DEVICE_MANAGER_DIAGNOSTICS = Diagnostics()
 69EVENT_DIAGNOSTICS = Diagnostics()
 70EVENT_MEDIA_DIAGNOSTICS = Diagnostics()
 71STREAMING_MANAGER_DIAGNOSTICS = Diagnostics()
 72
 73MAP = {
 74    "subscriber": SUBSCRIBER_DIAGNOSTICS,
 75    "device_manager": DEVICE_MANAGER_DIAGNOSTICS,
 76    "event": EVENT_DIAGNOSTICS,
 77    "event_media": EVENT_MEDIA_DIAGNOSTICS,
 78    "streaming_manager": STREAMING_MANAGER_DIAGNOSTICS,
 79}
 80
 81
 82def reset() -> None:
 83    """Clear all diagnostics, for testing."""
 84    for diagnostics in MAP.values():
 85        diagnostics.reset()
 86
 87
 88def get_diagnostics() -> dict[str, Any]:
 89    """Produce diagnostics information for the library."""
 90    return {k: v.as_dict() for (k, v) in MAP.items() if v.as_dict()}
 91
 92
 93REDACT_KEYS = {
 94    "name",
 95    "custom_name",
 96    "displayName",
 97    "parent",
 98    "assignee",
 99    "subject",
100    "object",
101    "userId",
102    "resourceGroup",
103    "eventId",
104    "eventSessionId",
105    "eventThreadId",
106}
107REDACTED = "**REDACTED**"
108
109
110T = TypeVar("T")
111
112
113def redact_data(data: T) -> T | dict | list:
114    """Redact sensitive data in a dict."""
115    if not isinstance(data, (Mapping, list)):
116        return data
117
118    if isinstance(data, list):
119        return cast(T, [redact_data(item) for item in data])
120
121    redacted = {**data}
122
123    for key, value in redacted.items():
124        if key in REDACT_KEYS:
125            redacted[key] = REDACTED
126        elif isinstance(value, dict):
127            redacted[key] = redact_data(value)
128        elif isinstance(value, list):
129            redacted[key] = [redact_data(item) for item in value]
130
131    return redacted