From 1ceb06fb1642f978ddc6b391b5f04f721f8f726e Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 9 May 2015 14:47:40 +0800
Subject: [PATCH] dds: use context manager for batches

---
 artiq/coredevice/dds.py                        | 16 ++++++++++++++++
 examples/master/repository/dds_test.py         |  8 ++++----
 examples/master/repository/photon_histogram.py |  7 +++----
 3 files changed, 23 insertions(+), 8 deletions(-)

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