Source code for semi_cr.core.station.instruments.qcodes_proxy

import logging
from abc import ABC
from collections.abc import Mapping
from dataclasses import dataclass, field
from io import StringIO
from typing import Any

import qcodes
from qcodes import Instrument
from ruamel.yaml import YAML

from semi_cr.core.station.instruments import LabInstrument

type QCoDeSInstrumentSpec = tuple[str, Mapping[str, Any]]




class QCoDeSStationSingleton:
    """We need a Station temporarily in order to construct the QCoDeS instruments from config. Maybe we could benefit
    from having an actual persistent QCoDeS station on NQCPStation and using that. I'm concerned about separation of
    concerns, so I'm opting for the smallest thing for now."""

    _station: qcodes.Station | None = None



    @classmethod
    def get_instance(cls) -> qcodes.Station:
        if cls._station is None:
            cls._station = qcodes.Station()
        return cls._station






@dataclass
class QCoDeSInstrumentProxy(LabInstrument, ABC):
    """Base class for static instruments backed by a live QCoDeS instrument. QCoDeS represents a connection to an
    instrument as the _presence_ of a driver object. I.e. if the object exists, we may expect there to be a
    connection (more or less)... We wrap that behavior by maintaining a copy of the specification from the config
    files, and only inflating the actual QCoDeS instrument when it's actually needed."""

    qcodes_config: QCoDeSInstrumentSpec = field(repr=False)
    _qcodes_instrument: Instrument | None = field(default=None, init=False, repr=False)



    def connect(self) -> None:
        """In QCoDeS, connection means instantiating the driver with the configured instrument."""
        if self._qcodes_instrument is not None:
            return
        name, config = self.qcodes_config
        stream = StringIO()
        YAML().dump(
            {
                "instruments": {
                    name: config,
                }
            },
            stream,
        )

        station = QCoDeSStationSingleton.get_instance()
        station.load_config(stream.getvalue())
        try:
            self._qcodes_instrument = station.load_instrument(
                name,
                update_snapshot=False,
            )
        except Exception:
            # Some drivers register their QCoDeS instance before their
            # constructor finishes. Remove that partial instance so Retry can
            # use the same configured name.
            try:
                components = getattr(station, "components", {})
                partial = components.get(name) if isinstance(components, Mapping) else None
                if partial is not None:
                    station.remove_component(name)
                    partial.close()
                elif Instrument.exist(name):
                    Instrument.find_instrument(name).close()
            except Exception:
                logging.exception("Could not clean up partially attached instrument %s", name)
            raise




    def is_connected(self):
        return self._qcodes_instrument is not None




    def disconnect(self):
        if self._qcodes_instrument is not None:
            instrument = self._qcodes_instrument
            try:
                instrument.close()
            finally:
                self._qcodes_instrument = None
                station = QCoDeSStationSingleton.get_instance()
                components = getattr(station, "components", {})
                if isinstance(components, Mapping) and components.get(self.name) is instrument:
                    station.remove_component(self.name)