RtlsdrAio improvements

.travis.yml

@@ -4,7 +4,7 @@ dist: trusty
 python:
   - "2.7"
   - "3.4"
-  - "3.5"
+  - "3.5.2"
 addons:
   apt:
     packages:

README.md

@@ -78,7 +78,7 @@ async def streaming():
   # ...
 
   # to stop streaming:
-  sdr.stop()
+  await sdr.stop()
 
  # done
  sdr.close()

rtlsdr/rtlsdraio.py

@@ -32,39 +32,53 @@ def __init__(self, func_start, func_stop=None, queue_size=20, *, loop=None):
 
         self.running = False
 
-    def _callback(self, *args):
-        if self.running and not self.queue.full():
-            self.loop.call_soon_threadsafe(self.queue.put_nowait, args)
-        else:
+    async def add_to_queue(self, *args):
+        try:
+            self.queue.put_nowait(args)
+        except asyncio.QueueFull:
             log.info('extra callback data lost')
 
-    def start(self):
+    def _callback(self, *args):
+        if not self.running:
+            return
+        asyncio.run_coroutine_threadsafe(self.add_to_queue(*args), self.loop)
+
+    async def start(self):
         assert(not self.running)
 
         # start legacy async function
         future = self.loop.run_in_executor(None, self.func_start, self._callback)
         asyncio.ensure_future(future, loop=self.loop)
+        self.executor_task = future
         self.running = True
 
-    def stop(self):
+    async def stop(self):
         assert(self.running)
 
-        # send a signal to stop
-        self.queue.put_nowait((StopAsyncIteration(),))
         self.running = False
 
+        # send a signal to stop
+        iter_stopped = False
+        while not iter_stopped:
+            try:
+                self.queue.put_nowait((StopAsyncIteration(),))
+                iter_stopped = True
+            except asyncio.QueueFull:
+                try:
+                    self.queue.task_done()
+                except ValueError:
+                    pass
         if self.func_stop:
             # stop legacy async function
-            future = self.loop.run_in_executor(None, self.func_stop)
-            asyncio.ensure_future(future, loop=self.loop)
-
-            #self.func_stop()
+            await self.loop.run_in_executor(None, self.func_stop)
+        await self.executor_task
 
     async def __aiter__(self):
         return self
 
     async def __anext__(self):
         val = await self.queue.get()
+        self.queue.task_done()
 
         if isinstance(val[0], StopAsyncIteration):
             raise StopAsyncIteration
@@ -107,10 +121,10 @@ def stream(self, num_samples_or_bytes=DEFAULT_READ_SIZE, format='samples', loop=
         self.async_iter = AsyncCallbackIter(func_start=lambda cb: func_start(cb, num_samples_or_bytes),
                                             func_stop=self.cancel_read_async,
                                             loop=loop)
-        self.async_iter.start()
+        asyncio.ensure_future(self.async_iter.start(), loop=loop)
 
         return self.async_iter
 
     def stop(self):
         ''' Stop async stream. '''
-        self.async_iter.stop()
+        return asyncio.ensure_future(self.async_iter.stop(), loop=self.async_iter.loop)

tests/test_aio.py

@@ -1,10 +1,11 @@
 import asyncio
+import pytest
 
 def test(rtlsdraio):
-    async def main():
-        import math
-        from utils import generic_test
+    import math
+    from utils import generic_test
 
+    async def main():
         sdr = rtlsdraio.RtlSdrAio()
         generic_test(sdr)
 
@@ -16,27 +17,39 @@ async def main():
         print('  center frequency %0.6f MHz' % (sdr.fc/1e6))
         print('  gain: %d dB' % sdr.gain)
 
+
         print('Streaming samples...')
+        await process_samples(sdr, 'samples')
+        await sdr.stop()
+
+        print('Streaming bytes...')
+        await process_samples(sdr, 'bytes')
+        await sdr.stop()
+
+        # make sure our format parameter checks work
+        with pytest.raises(ValueError):
+            await process_samples(sdr, 'foo')
+
+        print('Done')
+
+        sdr.close()
+
+
+    async def process_samples(sdr, fmt):
+        async def packed_bytes_to_iq(samples):
+            return sdr.packed_bytes_to_iq(samples)
 
         i = 0
-        async for samples in sdr.stream():
+        async for samples in sdr.stream(format=fmt):
+            if fmt == 'bytes':
+                samples = await packed_bytes_to_iq(samples)
             power = sum(abs(s)**2 for s in samples) / len(samples)
             print('Relative power:', 10*math.log10(power), 'dB')
 
             i += 1
 
             if i > 20:
-                sdr.stop()
                 break
 
-        print('Done')
-
-        sdr.close()
-
-    async def do_nothing():
-        for i in range(50):
-            await asyncio.sleep(0.1)
-            print('#')
-
     loop = asyncio.get_event_loop()
-    loop.run_until_complete(asyncio.wait([main(), do_nothing()]))
+    loop.run_until_complete(main())