dds: use context manager for batches

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

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:

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):