metrics.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Final, Iterable


@dataclass(frozen=True)
class RenderColumn:
    label: str
    metric: str
    avg_label: str | None = None


@dataclass(frozen=True)
class ShowSpec:
    flag: str
    metric_json_keys: dict[str, str]
    columns: tuple[RenderColumn, ...]


SHOW_SPECS: Final[tuple[ShowSpec, ...]] = (
    ShowSpec(
        flag="showtemp",
        metric_json_keys={"temp_c": "Temperature (Sensor core) (C)"},
        columns=(RenderColumn(label="Temp", metric="temp_c"),),
    ),
    ShowSpec(
        flag="showpower",
        metric_json_keys={"avg_pwr_w": "Average Graphics Package Power (W)"},
        columns=(RenderColumn(label="AvgPwr", metric="avg_pwr_w", avg_label="AvgPwr"),),
    ),
    ShowSpec(
        flag="showhcuclocks",
        metric_json_keys={"sclk_mhz": "sclk clock speed"},
        columns=(RenderColumn(label="sclk", metric="sclk_mhz"),),
    ),
    ShowSpec(
        flag="showmemuse",
        metric_json_keys={"vram_pct": "HCU memory use (%)"},
        columns=(RenderColumn(label="VRAM%", metric="vram_pct"),),
    ),
    ShowSpec(
        flag="showuse",
        metric_json_keys={"hcu_pct": "HCU use (%)"},
        columns=(RenderColumn(label="GPU%", metric="hcu_pct", avg_label="GPU%"),),
    ),
)

SPEC_BY_FLAG: Final[dict[str, ShowSpec]] = {spec.flag: spec for spec in SHOW_SPECS}

SUPPORTED_SHOW_FLAGS: Final[tuple[str, ...]] = tuple(spec.flag for spec in SHOW_SPECS)

DEFAULT_SHOW_FLAGS: Final[tuple[str, ...]] = SUPPORTED_SHOW_FLAGS

JSON_KEY_BY_METRIC: Final[dict[str, str]] = {
    metric: json_key
    for spec in SHOW_SPECS
    for metric, json_key in spec.metric_json_keys.items()
}


def normalized_show_flags(show_flags: Iterable[str] | None) -> list[str]:
    """Normalize user-selected show flags, preserving order and uniqueness."""

    if not show_flags:
        return list(DEFAULT_SHOW_FLAGS)
    output: list[str] = []
    for flag in show_flags:
        if flag not in SPEC_BY_FLAG:
            continue
        if flag not in output:
            output.append(flag)
    return output if output else list(DEFAULT_SHOW_FLAGS)


def hy_smi_args_for_show_flags(show_flags: Iterable[str], wait_idle: bool) -> list[str]:
    """Build hy-smi args with JSON output and requested fine-grained metrics."""

    ordered_flags = normalized_show_flags(show_flags)
    if wait_idle:
        # wait-idle relies on usage+memory metrics even when not displayed.
        for required in ("showmemuse", "showuse"):
            if required not in ordered_flags:
                ordered_flags.append(required)
    args = ["--json"]
    args.extend(f"--{flag}" for flag in ordered_flags)
    return args


def render_columns_for_show_flags(show_flags: Iterable[str]) -> list[RenderColumn]:
    """Resolve display columns from ordered show flags."""

    columns: list[RenderColumn] = []
    for flag in normalized_show_flags(show_flags):
        columns.extend(SPEC_BY_FLAG[flag].columns)
    return columns