Support fast discovery if one device is found

By passing abort_on_found when scanning, the process will abort even if the
timeout has not been reached. This is now default in atvremote and
pyatv.helpers.auto_connect. This fixes #19.

README.rst

@@ -135,7 +135,7 @@ Tasks related to library features
 - Non-polling based API (callbacks) (#8)
 - Send URL to AirPlay media **DONE** (#16)
 - Arrow keys (up, down, left and right) (#17)
-- Allow auto discovery stop after finding a device (#19)
+- Allow auto discovery stop after finding a device **DONE** (#19)
 - Better output for "playing" in atvremote (#20)
 - Verify compatibility with python > 3.5 (tox) *Pending* (#18)
 - Fix exit code in atvremote

docs/finding_devices.rst

@@ -59,3 +59,39 @@ an async call. A simple example might look like this:
     loop.run_until_complete(discover(loop))
 
 API Reference: :py:meth:`pyatv.scan_for_apple_tvs`
+
+Finding a single device
+-----------------------
+Under some circumstance you might not care about which device you connect to,
+usually when you only have one device on the network. To simplify and speed up
+the discovery process, you can set the flag ``abort_on_found`` to ``True``.
+This will make ``pyatv.scan_for_apple_tvs`` abort when a device has been found,
+thus ignore the timeout and return quicker:
+
+.. code:: python
+
+    atvs = yield from pyatv.scan_for_apple_tvs(
+        loop, timeout=5, abort_on_found=True)
+
+This is for instance default behavior when using ``-a`` with atvremote. There's
+also a helper method that utilizes this by default:
+
+.. code:: python
+
+    import asyncio
+    from pyatv import helpers
+
+    @asyncio.coroutine
+    def print_what_is_playing(atv):
+        playing = yield from atv.metadata.playing()
+        print('Currently playing: ')
+        print(playing)
+
+    helpers.auto_connect(print_what_is_playing)
+
+When writing simpler application, ``auto_connect`` can be quite convenient as
+it can handle loop management for you. It is also possible to pass an error
+handler, that is called when a device is not found. See the API referece for
+more details.
+
+API Reference: :py:meth:`pyatv.helpers.auto_connect`

pyatv/__init__.py

@@ -2,8 +2,9 @@
 
 import asyncio
 import logging
-
+import concurrent
 import ipaddress
+
 from collections import namedtuple
 from zeroconf import ServiceBrowser, Zeroconf
 from aiohttp import ClientSession
@@ -32,33 +33,52 @@ def __new__(cls, name, address, login_id, port=3689):
 # pylint: disable=too-few-public-methods
 class _ServiceListener(object):
 
-    def __init__(self):
+    def __init__(self, abort_on_found, semaphore):
         """Initialize a new _ServiceListener."""
+        self.abort_on_found = abort_on_found
+        self.semaphore = semaphore
         self.found_devices = []
 
     def add_service(self, zeroconf, service_type, name):
         """Callback from zeroconf when a service has been discovered."""
+        # If it's not instantly lockable, then abort_on_found is True and we
+        # have found a device already
+        if not self.semaphore.locked():
+            return
+
         info = zeroconf.get_service_info(service_type, name)
         if info.type == HOMESHARING_SERVICE:
-            address = ipaddress.ip_address(info.address)
-            tv_name = info.properties[b'Name'].decode('utf-8')
-            hsgid = info.properties[b'hG'].decode('utf-8')
-            self.found_devices.append(AppleTVDevice(tv_name, address, hsgid))
-            _LOGGER.debug('Auto-discovered service %s at %s (hsgid: %s)',
-                          tv_name, address, hsgid)
+            self.add_device(info)
+
+            # Check if we should continue to run or not
+            if self.abort_on_found:
+                _LOGGER.debug('Aborting since a device was found')
+                self.semaphore.release()
         else:
             _LOGGER.warning('Discovered unknown device: %s', info)
 
+    def add_device(self, info):
+        """Add a new device to discovered list."""
+        address = ipaddress.ip_address(info.address)
+        tv_name = info.properties[b'Name'].decode('utf-8')
+        hsgid = info.properties[b'hG'].decode('utf-8')
+        self.found_devices.append(AppleTVDevice(tv_name, address, hsgid))
+        _LOGGER.debug('Auto-discovered service %s at %s (hsgid: %s)',
+                      tv_name, address, hsgid)
+
 
 @asyncio.coroutine
-def scan_for_apple_tvs(loop, timeout=5):
+def scan_for_apple_tvs(loop, timeout=5, abort_on_found=False):
     """Scan for Apple TVs using zeroconf (bonjour) and returns them."""
-    listener = _ServiceListener()
+    semaphore = asyncio.Semaphore(value=0, loop=loop)
+    listener = _ServiceListener(abort_on_found, semaphore)
     zeroconf = Zeroconf()
     try:
         ServiceBrowser(zeroconf, HOMESHARING_SERVICE, listener)
         _LOGGER.debug('Discovering devices for %d seconds', timeout)
-        yield from asyncio.sleep(timeout, loop=loop)
+        yield from asyncio.wait_for(semaphore.acquire(), timeout, loop=loop)
+    except concurrent.futures.TimeoutError:
+        pass  # Will happen when timeout occurs (totally normal)
     finally:
         zeroconf.close()
 

pyatv/__main__.py

@@ -121,8 +121,8 @@ def _print_found_apple_tvs(atvs, outstream=sys.stdout):
 
 @asyncio.coroutine
 def _handle_autodiscover(args, loop):
-    atvs = yield from pyatv.scan_for_apple_tvs(loop,
-                                               timeout=args.scan_timeout)
+    atvs = yield from pyatv.scan_for_apple_tvs(
+        loop, timeout=args.scan_timeout, abort_on_found=True)
     if len(atvs) == 0:
         logging.error('Could not find any Apple TV on current network')
         return 1

pyatv/helpers.py

@@ -4,7 +4,7 @@
 import pyatv
 
 
-def auto_connect(handler, not_found=None, event_loop=None):
+def auto_connect(handler, timeout=5, not_found=None, event_loop=None):
     """Convenient method for connecting to a device.
 
     This is a convenience method that create an event loop, auto discovers
@@ -20,7 +20,8 @@ def auto_connect(handler, not_found=None, event_loop=None):
     # the event loop
     @asyncio.coroutine
     def _handle(loop):
-        atvs = yield from pyatv.scan_for_apple_tvs(loop, timeout=5)
+        atvs = yield from pyatv.scan_for_apple_tvs(
+            loop, timeout=timeout, abort_on_found=True)
 
         # Take the first device found
         if len(atvs) > 0:

tests/test_functional.py

@@ -26,6 +26,11 @@
 # (extracted form a real device)
 PAIRINGCODE = '690E6FF61E0D7C747654A42AED17047D'
 
+HOMESHARING_SERVICE_1 = zeroconf_stub.homesharing_service(
+    'AAAA', b'Apple TV 1', '10.0.0.1', b'aaaa')
+HOMESHARING_SERVICE_2 = zeroconf_stub.homesharing_service(
+    'BBBB', b'Apple TV 2', '10.0.0.2', b'bbbb')
+
 
 class FunctionalTest(AioHTTPTestCase):
 
@@ -60,16 +65,24 @@ def get_connected_device(self, identifier):
 
     @unittest_run_loop
     def test_scan_for_apple_tvs(self):
-        zeroconf_stub.stub(pyatv, zeroconf_stub.homesharing_service(
-            'AAAA', b'Apple TV', '10.0.0.1', b'aaaa'))
+        zeroconf_stub.stub(pyatv, HOMESHARING_SERVICE_1)
 
         atvs = yield from pyatv.scan_for_apple_tvs(self.loop, timeout=0)
         self.assertEqual(len(atvs), 1)
-        self.assertEqual(atvs[0].name, 'Apple TV')
+        self.assertEqual(atvs[0].name, 'Apple TV 1')
         self.assertEqual(atvs[0].address, ipaddress.ip_address('10.0.0.1'))
         self.assertEqual(atvs[0].login_id, 'aaaa')
         self.assertEqual(atvs[0].port, 3689)
 
+    @unittest_run_loop
+    def test_scan_abort_on_first_found(self):
+        zeroconf_stub.stub(pyatv, HOMESHARING_SERVICE_1, HOMESHARING_SERVICE_2)
+
+        atvs = yield from pyatv.scan_for_apple_tvs(
+            self.loop, timeout=0, abort_on_found=True)
+        self.assertEqual(len(atvs), 1)
+        self.assertEqual(atvs[0].name, 'Apple TV 1')
+
     # This is not a pretty test and it does crazy things. Should probably be
     # re-written later but will do for now.
     @unittest_run_loop