Move pushback buffer into vm so that it will work with other kinds of stream; move stream closing into same for neatness; support cancelling capturing through asyncio.Task cancellation

scope.py

@@ -67,14 +67,6 @@ class Scope(vm.VirtualMachine):
         self._awg_running = False
         Log.info(f"Initialised scope, revision: {revision}")
 
-    def close(self):
-        if self._writer is not None:
-            self._writer.close()
-            self._writer = None
-            self._reader = None
-
-    __del__ = close
-
     def calculate_lo_hi(self, low, high, params=None):
         params = self.analog_params if params is None else self.AnalogParams(*params)
         l = (low - params.offset) / params.scale
@@ -212,15 +204,18 @@ class Scope(vm.VirtualMachine):
             await self.issue_configure_device_hardware()
             await self.issue_triggered_trace()
         while True:
+            try:
                 code, timestamp = (int(x, 16) for x in await self.read_replies(2))
-            if code != 2:
+                if code != vm.TraceStatus.Wait:
                     break
+            except asyncio.CancelledError:
+                await self.issue_cancel_trace()
         start_timestamp = timestamp - nsamples*ticks*clock_scale
         if start_timestamp < 0:
             start_timestamp += 1<<32
             timestamp += 1<<32
         address = int((await self.read_replies(1))[0], 16)
-        if capture_mode.BufferMode in {vm.BufferMode.Chop, vm.BufferMode.MacroChop, vm.BufferMode.ChopDual}:
+        if capture_mode.analog_channels == 2:
             address -= address % 2
         traces = DotDict()
         for dump_channel, channel in enumerate(sorted(analog_channels)):
@@ -412,7 +407,12 @@ async def main():
     s = await Scope.connect(args.device)
 
 def await(g):
-    return asyncio.get_event_loop().run_until_complete(g)
+    task = asyncio.Task(g)
+    while True:
+        try:
+            return asyncio.get_event_loop().run_until_complete(task)
+        except KeyboardInterrupt:
+            task.cancel()
 
 def capture(*args, **kwargs):
     return await(s.capture(*args, **kwargs))

streams.py

@@ -56,7 +56,6 @@ class SerialStream:
         self._loop = loop if loop is not None else asyncio.get_event_loop()
         self._output_buffer = bytes()
         self._output_buffer_empty = None
-        self._pushback_buffer = bytes()
 
     def __repr__(self):
         return f'<{self.__class__.__name__}:{self._device}>'
@@ -105,14 +104,7 @@ class SerialStream:
             self._output_buffer_empty.set_result(None)
             self._output_buffer_empty = None
 
-    def pushback(self, data):
-        self._pushback_buffer += data
-
     async def read(self, n=None):
-        if self._pushback_buffer:
-            data = self._pushback_buffer if n is None else self._pushback_buffer[:n]
-            self._pushback_buffer = self._pushback_buffer[len(data):]
-            return data
         while True:
             w = self._connection.in_waiting
             if w:

vm.py

@@ -151,6 +151,12 @@ class KitchenSinkB(IntEnum):
     AnalogFilterEnable       = 0x80
     WaveformGeneratorEnable  = 0x40
 
+class TraceStatus(IntEnum):
+    Done = 0x00
+    Auto = 0x01
+    Wait = 0x02
+    Stop = 0x03
+
 CaptureMode = namedtuple('CaptureMode', ('clock_low', 'clock_high', 'clock_max', 'analog_channels', 'sample_width',
                                          'logic_channels', 'clock_divide', 'TraceMode', 'BufferMode'))
 
@@ -226,7 +232,8 @@ class VirtualMachine:
             self._vm._transactions.append(self)
             return self
         async def __aexit__(self, exc_type, exc_value, traceback):
-            self._vm._transactions.pop()
+            if self._vm._transactions.pop() != self:
+                raise RuntimeError("Mis-ordered transactions")
             if exc_type is None:
                 await self._vm.issue(self._data)
             return False
@@ -235,6 +242,15 @@ class VirtualMachine:
         self._reader = reader
         self._writer = writer
         self._transactions = []
+        self._reply_buffer = b''
+
+    def close(self):
+        if self._writer is not None:
+            self._writer.close()
+            self._writer = None
+            self._reader = None
+
+    __del__ = close
 
     def transaction(self):
         return self.Transaction(self)
@@ -256,7 +272,7 @@ class VirtualMachine:
         if self._transactions:
             raise TypeError("Command transaction in progress")
         replies = []
-        data = b''
+        data, self._reply_buffer = self._reply_buffer, b''
         while len(replies) < n:
             index = data.find(b'\r')
             if index >= 0:
@@ -267,7 +283,7 @@ class VirtualMachine:
             else:
                 data += await self._reader.read()
         if data:
-            self._reader.pushback(data)
+            self._reply_buffer = data
         return replies
 
     async def reset(self):
@@ -278,6 +294,7 @@ class VirtualMachine:
         await self._writer.drain()
         while not (await self._reader.read()).endswith(b'!'):
             pass
+        self._reply_buffer = b''
         Log.debug("Reset complete")
 
     async def set_registers(self, **kwargs):