venv/lib/python3.11/site-packages/wandb/plot/viz.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from wandb.data_types import Table
from wandb.errors import Error


@dataclass
class VisualizeSpec:
    name: str
    key: str = ""

    @property
    def config_value(self) -> dict[str, Any]:
        return {
            "id": self.name,
            "historyFieldSettings": {"x-axis": "_step", "key": self.key},
        }

    @property
    def config_key(self) -> tuple[str, str, str]:
        return ("_wandb", "viz", self.key)


@dataclass
class Visualize:
    table: Table
    spec: VisualizeSpec

    def set_key(self, key: str) -> None:
        self.spec.key = key


def visualize(id: str, value: Table) -> Visualize:
    if not isinstance(value, Table):
        raise Error(
            f"Expected `value` to be `wandb.Table` type, instead got {type(value).__name__}"
        )
    return Visualize(table=value, spec=VisualizeSpec(name=id))