InstroAWG("name", driver=RigolDG1022Z(visa_resource="USB0::..."), num_channels=2)

from instro.unstable.awg.drivers import RigolDG1022Z
from instro.unstable.awg import InstroAWG

awg = InstroAWG(
    name="myAWG",
    driver=RigolDG1022Z(visa_resource="USB0::0x1AB1::0x0642::DG1ZA000000000::INSTR"),
    num_channels=2,
)

from instro.lib.transports import VisaDriver

visa = VisaDriver("USB0::...")
try:
    visa.open()
    print(visa.query("*IDN?"))
finally:
    visa.close()

from instro.unstable.awg.drivers import RigolDG1022Z
from instro.unstable.awg import InstroAWG
from instro.unstable.awg.types import AmplitudeMeasurementUnit, Sine
from instro.lib.publishers import NominalCorePublisher

VISA_RESOURCE = "USB0::0x1AB1::0x0642::DG1ZA000000000::INSTR"
DATASET_RID = "<your dataset here>"

awg = InstroAWG(
    name="myAWG",
    driver=RigolDG1022Z(visa_resource=VISA_RESOURCE),
    num_channels=2,
)
awg.add_publisher(NominalCorePublisher(dataset_rid=DATASET_RID))

awg.open()

# Configure channel 1
awg.set_waveform(1, Sine(frequency_hz=1000.0))
awg.set_amplitude(1, 2.0, AmplitudeMeasurementUnit.VPP)
awg.output_enable(1, True)

# Read channel 1
waveform = awg.get_waveform(1)
enabled = awg.get_output_state(1)
print(f"Waveform: {waveform}, Output enabled: {enabled.latest}")

awg.output_enable(1, False)
awg.close()

awg = InstroAWG(
    name="labAWG",
    driver=MyVendorAWG(host="10.0.0.42"),
    num_channels=1,
)

def open(self) -> None:
    """Open the underlying transport."""

def close(self) -> None:
    """Close the underlying transport."""

def set_waveform(self, channel: int, waveform: Waveform) -> None:
    """Program channel with the waveform definition; raise ValueError if the definition is unsupported."""

def get_waveform(self, channel: int) -> Waveform:
    """Get the current waveform on channel; drivers may return the last-programmed definition if not readable."""

def set_amplitude(self, channel: int, amplitude: float, unit: AmplitudeMeasurementUnit) -> None:
    """Set the output amplitude on channel."""

def get_amplitude(self, channel: int) -> tuple[float, AmplitudeMeasurementUnit]:
    """Get the current output amplitude and voltage unit on channel."""

def set_offset(self, channel: int, offset: float) -> None:
    """Set the DC offset (volts) on channel."""

def get_offset(self, channel: int) -> float:
    """Get the DC offset (volts) on channel."""

def output_enable(self, channel: int, enable: bool) -> None:
    """Enable or disable the output on channel."""

def get_output_state(self, channel: int) -> bool:
    """Return True if the output on channel is enabled."""

from instro.unstable.awg.awg import AWGDriverBase
from instro.unstable.awg.types import AmplitudeMeasurementUnit, Sine, Waveform
from instro.lib.transports import VisaConfig, VisaDriver


class RigolDG1022Z(AWGDriverBase):
    """SCPI driver for the Rigol DG1022Z two-channel arbitrary waveform generator."""

    def __init__(self, visa_resource: str | VisaConfig) -> None:
        self._visa = VisaDriver(visa_resource)

    def open(self) -> None:
        self._visa.open()

    def close(self) -> None:
        self._visa.close()

    def set_waveform(self, channel: int, waveform: Waveform) -> None:
        with self._visa.lock():
            if isinstance(waveform, Sine):
                self._visa.write(f":SOUR{channel}:FUNC SIN")
                self._visa.write(f":SOUR{channel}:FREQ {waveform.frequency_hz}")
                self._visa.write(f":SOUR{channel}:PHAS {waveform.phase_deg % 360.0}")
            # ... other Waveform shapes ...
            else:
                raise ValueError(f"unsupported waveform definition {type(waveform).__name__}")
            self._check_errors()  # one check per call, not one per write — see note above

    def set_amplitude(self, channel: int, amplitude: float, unit: AmplitudeMeasurementUnit) -> None:
        if unit is AmplitudeMeasurementUnit.VP:
            raise ValueError("the DG1022Z has no VP amplitude unit; convert to VPP, VRMS, or DBM")
        with self._visa.lock():
            self._visa.write(f":SOUR{channel}:VOLT:UNIT {unit.value}")
            self._visa.write(f":SOUR{channel}:VOLT {amplitude}")
            self._check_errors()

    # ... other required and optional methods ...

    def _check_errors(self) -> None:
        err = self._visa.query(":SYST:ERR?")
        code = err.strip().split(",", 1)[0].lstrip("+")
        if code != "0":
            raise RuntimeError(f"Rigol DG1022Z reported error: {err.strip()}")

from instro.unstable.awg import InstroAWG
from instro.unstable.awg.awg import AWGDriverBase
from instro.unstable.awg.types import AmplitudeMeasurementUnit, Waveform, Sine
from instro.lib.transports import VisaDriver


class MyCustomAWGDriver(AWGDriverBase):
    """Custom driver for my lab's proprietary AWG."""

    def __init__(self, visa_resource: str) -> None:
        self._visa = VisaDriver(visa_resource)

    def open(self) -> None:
        self._visa.open()

    def close(self) -> None:
        self._visa.close()

    def set_waveform(self, channel: int, waveform: Waveform) -> None:
        if isinstance(waveform, Sine):
            self._visa.write(f"CH{channel}:WAVE SIN,{waveform.frequency_hz}")
            self._check_errors()
        else:
            raise ValueError(f"unsupported waveform definition {type(waveform).__name__}")

    # ... implement other required methods ...

    def get_output_state(self, channel: int) -> bool:
        return self._visa.query(f"CH{channel}:OUTP?") == "ON"

    def _check_errors(self) -> None:
        err = self._visa.query("ERR?")
        if err != "OK":
            raise RuntimeError(f"AWG error: {err}")


awg = InstroAWG(
    name="labAWG",
    driver=MyCustomAWGDriver(visa_resource="<VISA_ADDRESS>"),
    num_channels=1,
)

awg.open()
awg.set_waveform(1, Sine(frequency_hz=1000.0))
awg.close()