Source code for semi_cr.core.lab.station.context.static

from dataclasses import dataclass, field
from typing import Literal

import networkx as nx

from semi_cr.core.lab.station.base import (
    ConfigBundle,
    InstrumentSpec,
)
from semi_cr.core.lab.station.graphing.context import ConfigContext
from semi_cr.core.lab.station.graphing.static_builder import (
    build_wiring_graph,
)
from semi_cr.core.lab.station.graphing.view import graph_view

from .config import build_instrument_specs




@dataclass
class InfrastructureContext:
    wiring_config: ConfigBundle
    electronics_config: ConfigBundle

    config: ConfigBundle = field(init=False)
    graph: nx.MultiDiGraph = field(init=False)

    def __post_init__(self) -> None:
        self.config = ConfigBundle.combine(
            self.wiring_config,
            self.electronics_config,
        )
        self.graph = self.build_graph()



    def build_graph(self) -> nx.MultiDiGraph:
        return build_wiring_graph(self.config.config)




    def view_graph(
        self,
        what: Literal["nodes", "edges", "both"] = "both",
        include_attrs: bool = True,
    ):
        return graph_view(
            self.graph,
            what=what,
            include_attrs=include_attrs,
        )






@dataclass
class StaticStationContext:
    """
    Description of the complete station without talking to hardware.
    """
    device: ConfigContext
    infrastructure: InfrastructureContext

    config: ConfigBundle = field(init=False)
    instruments: dict[str, InstrumentSpec] = field(init=False)
    graph: nx.MultiDiGraph = field(init=False)

    def __post_init__(self) -> None:
        self.config = ConfigBundle.combine(
            self.device.config,
            self.infrastructure.config,
        )
        self.instruments = build_instrument_specs(
            self.config.config
        )
        self.graph = build_wiring_graph(
            self.config.config
        )