Source code for semi_cr.core.lab.station.snapshot

from __future__ import annotations

from datetime import datetime
from pathlib import Path
from typing import Any

import yaml
from pydantic import BaseModel, Field




class ParameterSnapshot(BaseModel):
    value: Any = None
    unit: str | None = None
    label: str | None = None





class ComponentSnapshot(BaseModel):
    name: str
    class_name: str | None = None
    label: str | None = None

    parameters: dict[str, ParameterSnapshot] = Field(default_factory=dict)
    submodules: dict[str, "ComponentSnapshot"] = Field(default_factory=dict)





class StationSnapshot(BaseModel):
    timestamp: datetime = Field(default_factory=datetime.now)
    instruments: dict[str, ComponentSnapshot] = Field(default_factory=dict)
    config: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)



    def to_yaml(self, path: str | Path) -> None:

        path = Path(path)

        data = self.model_dump(mode="json",exclude_none=True)

        with path.open("w", encoding="utf-8") as file:
            yaml.safe_dump(data, file,sort_keys=False)





def component_from_qcodes_snapshot(
    name: str,
    snapshot: dict[str, Any],
) -> ComponentSnapshot:

    parameters = {
        parameter_name: ParameterSnapshot(
            value=parameter.get("value"),
            unit=parameter.get("unit"),
            label=parameter.get("label"),
        )
        for parameter_name, parameter
        in snapshot.get("parameters", {}).items()
    }

    submodules = {
        submodule_name: component_from_qcodes_snapshot(
            submodule_name,
            submodule_snapshot,
        )
        for submodule_name, submodule_snapshot
        in snapshot.get("submodules", {}).items()
    }

    return ComponentSnapshot(
        name=snapshot.get("name", name),
        label=snapshot.get("label"),
        class_name=snapshot.get("__class__"),
        parameters=parameters,
        submodules=submodules,
    )




def station_snapshot_from_qcodes(
    station,
) -> StationSnapshot:

    raw = station.snapshot()

    instruments = {
        instrument_name: component_from_qcodes_snapshot(
            instrument_name,
            instrument_snapshot,
        )
        for instrument_name, instrument_snapshot
        in raw.get("instruments", {}).items()
    }

    return StationSnapshot(
        instruments=instruments,
        config=raw.get("config", {}),
    )