diff --git a/artiq/coredevice/dds.py b/artiq/coredevice/dds.py index 890ff9e96..7799c1d31 100644 --- a/artiq/coredevice/dds.py +++ b/artiq/coredevice/dds.py @@ -10,6 +10,19 @@ PHASE_MODE_ABSOLUTE = 1 PHASE_MODE_TRACKING = 2 +class _BatchContextManager: + def __init__(self, dds_bus): + self.dds_bus = dds_bus + + @kernel + def __enter__(self): + self.dds_bus.batch_enter() + + @kernel + def __exit__(self, type, value, traceback): + self.dds_bus.batch_exit() + + class DDSBus(AutoDB): """Core device Direct Digital Synthesis (DDS) bus batching driver. @@ -17,6 +30,9 @@ class DDSBus(AutoDB): class DBKeys: core = Device() + def build(self): + self.batch = _BatchContextManager(self) + @kernel def batch_enter(self): """Starts a DDS command batch. All DDS commands are buffered diff --git a/examples/master/repository/dds_test.py b/examples/master/repository/dds_test.py index 2c6e8e1d0..a36af8447 100644 --- a/examples/master/repository/dds_test.py +++ b/examples/master/repository/dds_test.py @@ -17,10 +17,10 @@ class DDSTest(Experiment, AutoDB): @kernel def run(self): - self.dds_bus.batch_enter() - self.dds1.set(120*MHz) - self.dds2.set(200*MHz) - self.dds_bus.batch_exit() + with self.dds_bus.batch: + self.dds1.set(120*MHz) + self.dds2.set(200*MHz) + delay(1*us) for i in range(10000): if i & 0x200: diff --git a/examples/master/repository/photon_histogram.py b/examples/master/repository/photon_histogram.py index 66ea3b5b7..aebe1beb4 100644 --- a/examples/master/repository/photon_histogram.py +++ b/examples/master/repository/photon_histogram.py @@ -27,10 +27,9 @@ class PhotonHistogram(Experiment, AutoDB): @kernel def program_cooling(self): - self.dds_bus.batch_enter() - self.bd_dds.set(200*MHz) - self.bdd_dds.set(300*MHz) - self.dds_bus.batch_exit() + with self.dds_bus.batch: + self.bd_dds.set(200*MHz) + self.bdd_dds.set(300*MHz) @kernel def cool_detect(self):