Source code for semi_cr.core.station.hydration

"""Hydration module for the NQCP station configuration. Hydration is what turns our configuration into more
"live"-like objects, such as Chip and Connector. It is akin to the QCoDeS way of creating a rich object from the
configuration, only here we declare a _handler_ per type, treating the types as regular strings instead of the QCoDeS
way of pointing to some class. That allows the two approaches to coexist."""

from __future__ import annotations

import logging
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

from semi_cr.core.station.chip import Chip
from semi_cr.core.station.graph import Connector
from semi_cr.core.station.instruments import LabInstrument
from semi_cr.core.station.instruments.dmm import DMMBlueprint
from semi_cr.core.station.instruments.opx1000 import OPX1000Blueprint
from semi_cr.core.station.instruments.spirack import SPIRackBlueprint

if TYPE_CHECKING:
    from semi_cr.core.lab.station.base import ConfigBundle
    from semi_cr.core.station.station import NQCPStation

CHIP_TYPE = "semi_cr.core.lab.instrument.chip.Chip"
CONNECTOR_TYPE = "semi_cr.core.lab.instrument.connector.Connector"
OPX1000_TYPE = "semi_cr.core.lab.drivers.opx_config.OPX1000"
SPIRACK_TYPE = "semi_cr.core.lab.drivers.spirack.QcodesSPIRack"
DMM_TYPE = "semi_cr.core.lab.instrument.routable_dmm.QcodesDMM"

InstrumentSpec = tuple[str, Mapping[str, Any]]
StaticLoader = Callable[[InstrumentSpec, Any], object]




@dataclass(frozen=True)
class UnloadedSpec:
    name: str
    type: str
    config: Mapping[str, Any]





@dataclass
class HydratedStationObjects:
    chips: list[Chip] = field(default_factory=list)
    connectors: list[Connector] = field(default_factory=list)
    instruments: list[LabInstrument] = field(default_factory=list)
    unloaded_specs: list[UnloadedSpec] = field(default_factory=list)





class MissingStaticLoaderError(Exception):
    pass





def load_chip(spec: InstrumentSpec, station: NQCPStation) -> Chip:
    return Chip.from_instrument_spec(spec, station)





def load_connector(spec: InstrumentSpec, station: NQCPStation) -> Connector:
    return Connector.from_instrument_spec(spec)





def load_opx1000(spec: InstrumentSpec, station: NQCPStation) -> OPX1000Blueprint:
    return OPX1000Blueprint.from_instrument_spec(spec)





def load_spirack(spec: InstrumentSpec, station: NQCPStation) -> SPIRackBlueprint:
    return SPIRackBlueprint.from_instrument_spec(spec)





def load_dmm(spec: InstrumentSpec, station: NQCPStation) -> DMMBlueprint:
    return DMMBlueprint.from_instrument_spec(spec)



STATIC_LOADERS: Mapping[str, StaticLoader] = {
    CHIP_TYPE: load_chip,
    CONNECTOR_TYPE: load_connector,
    OPX1000_TYPE: load_opx1000,
    SPIRACK_TYPE: load_spirack,
    DMM_TYPE: load_dmm,
}




def hydrate_station_config(config_bundle: ConfigBundle, station: NQCPStation) -> HydratedStationObjects:
    hydrated = HydratedStationObjects()

    for spec in config_bundle.get_items():
        name, config = spec
        try:
            loaded = hydrate_instrument_spec(spec, station)
        except MissingStaticLoaderError as error:
            spec_type = str(config.get("type", "<missing>"))
            logging.error("%s", error)
            hydrated.unloaded_specs.append(
                UnloadedSpec(
                    name=name,
                    type=spec_type,
                    config=config,
                )
            )
            continue

        if isinstance(loaded, Chip):
            hydrated.chips.append(loaded)
        elif isinstance(loaded, Connector):
            hydrated.connectors.append(loaded)
        elif isinstance(loaded, LabInstrument):
            hydrated.instruments.append(loaded)

    return hydrated





def hydrate_instrument_spec(spec: InstrumentSpec, station: NQCPStation) -> object:
    name, config = spec
    spec_type = config.get("type", "")
    loader = STATIC_LOADERS.get(spec_type)

    if loader is None:
        raise MissingStaticLoaderError(f"No static loader for instrument {name!r} with type {spec_type!r}")

    return loader(spec, station)