Source code for semi_cr.core.lab.pins.PinChannelList

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, TypeVar

from qcodes.instrument import ChannelList

from semi_cr.core.lab.pins.pins import (
    BasePinChannel,
    Drain,
    Gate,
    Ohmic,
    DCPinChannel,
    RFInput,
    RFOutput,
    RFPinChannel,
    Source,
)

if TYPE_CHECKING:
    from semi_cr.core.lab.devices.rf_device import BaseRFDevice
    from semi_cr.core.lab.devices.semiconductor import BaseSemiConductingDevice

C = TypeVar("C", bound=DCPinChannel)
D = TypeVar("D", bound=RFPinChannel)




@dataclass
class BasePinChannelList(ChannelList):

    def __init__(self, parent, name: str) -> None:

        super().__init__(
            parent,
            name=name,
            chan_type=BasePinChannel,
            snapshotable=True,
        )
    pass





class PinChannelList(BasePinChannelList):

    def __post_init__(self):

        try:
            super().__post_init__()
        except AttributeError:
            pass

        self.parent = BaseSemiConductingDevice
        self.chan_type = DCPinChannel



    def get_pin_by_name_and_type(self, pin_name: str, channel_type: type[C]) -> C:
        channel = self.get_channel_by_name(pin_name)
        if isinstance(channel, channel_type):
            return channel
        raise AttributeError("...")




    def get_source(self, source_name: str) -> Source:
        return self.get_pin_by_name_and_type(source_name, Source)




    def get_gate(self, gate_name: str) -> Gate:
        return self.get_pin_by_name_and_type(gate_name, Gate)




    def get_drain(self, drain_name: str) -> Drain:
        return self.get_pin_by_name_and_type(drain_name, Drain)




    def get_ohmic(self, ohmic_name: str) -> Ohmic:
        return self.get_pin_by_name_and_type(ohmic_name, Ohmic)


    @property
    def gates(self) -> list[Gate]:
        return [pin for pin in self._channels if isinstance(pin, Gate)]

    @property
    def sources(self) -> list[Source]:
        return [pin for pin in self._channels if isinstance(pin, Source)]

    @property
    def drains(self) -> list[Drain]:
        return [pin for pin in self._channels if isinstance(pin, Drain)]

    @property
    def ohmics(self) -> list[Ohmic]:
        return [pin for pin in self._channels if isinstance(pin, Ohmic)]





class RFPinChannelList(BasePinChannelList):

    def __post_init__(self):

        try:
            super().__post_init__()
        except AttributeError:
            pass

        self.parent = BaseRFDevice
        self.chan_type = RFPinChannel



    def get_pin_by_name_and_type(self, pin_name: str, channel_type: type[D]) -> D:
        channel = self.get_channel_by_name(pin_name)
        if isinstance(channel, channel_type):
            return channel
        raise AttributeError("...")




    def get_rf_input(self, source_name: str) -> RFInput:
        return self.get_pin_by_name_and_type(source_name, RFInput)




    def get_rf_output(self, gate_name: str) -> Gate:
        return self.get_pin_by_name_and_type(gate_name, RFOutput)


    @property
    def rf_inputs(self) -> list[RFInput]:
        return [pin for pin in self._channels if isinstance(pin, RFInput)]

    @property
    def rf_outputs(self) -> list[RFOutput]:
        return [pin for pin in self._channels if isinstance(pin, RFOutput)]