From 4bc845c53c26e38cfbfda2a528c62a5a7cdafbfa Mon Sep 17 00:00:00 2001
From: ViViDboarder <ViViDboarder@gmail.com>
Date: Wed, 24 Aug 2022 19:27:50 -0700
Subject: [PATCH] Allows more flexible scanning

This patch uses the BleakScanner interface to simplify the scanning as
well as provide more means of connecting. With an address, a BLEDevice,
or with a user provided scanner.

This change will allow me to implement a Home Assistant integration
using [Bluetooth
discovery](https://developers.home-assistant.io/docs/network_discovery/#best-practices-for-library-authors).

Other changes: `black`
---
 pyhatchbabyrest/constants.py            |  1 +
 pyhatchbabyrest/pyhatchbabyrestasync.py | 56 ++++++++++++++++---------
 2 files changed, 37 insertions(+), 20 deletions(-)

diff --git a/pyhatchbabyrest/constants.py b/pyhatchbabyrest/constants.py
index 447c462..99d633a 100644
--- a/pyhatchbabyrest/constants.py
+++ b/pyhatchbabyrest/constants.py
@@ -3,6 +3,7 @@
 COLOR_GRADIENT = (254, 254, 254)  # setting this color turns on Gradient mode
 CHAR_TX = "02240002-5efd-47eb-9c1a-de53f7a2b232"
 CHAR_FEEDBACK = "02260002-5efd-47eb-9c1a-de53f7a2b232"
+BT_MANUFACTURER_ID = 1076
 
 
 class PyHatchBabyRestSound(IntEnum):
diff --git a/pyhatchbabyrest/pyhatchbabyrestasync.py b/pyhatchbabyrest/pyhatchbabyrestasync.py
index 7fe2087..0ba871c 100644
--- a/pyhatchbabyrest/pyhatchbabyrestasync.py
+++ b/pyhatchbabyrest/pyhatchbabyrestasync.py
@@ -1,37 +1,53 @@
 import asyncio
+from typing import Optional
+from typing import Union
 
-from bleak import BleakClient, discover  # type: ignore
+from bleak import BleakClient
+from bleak import BleakScanner
+from bleak.backends.device import BLEDevice
 
-from .constants import CHAR_TX, CHAR_FEEDBACK, PyHatchBabyRestSound
+from .constants import BT_MANUFACTURER_ID
+from .constants import CHAR_FEEDBACK
+from .constants import CHAR_TX
+from .constants import PyHatchBabyRestSound
 
 
 class PyHatchBabyRestAsync(object):
-    """ An asynchronous interface to a Hatch Baby Rest device using bleak. """
-    def __init__(self, addr: str = None):
+    """An asynchronous interface to a Hatch Baby Rest device using bleak."""
+
+    def __init__(
+        self,
+        address_or_ble_device: Union[str, BLEDevice, None],
+        scanner: Optional[BleakScanner],
+    ):
         loop = asyncio.get_event_loop()
-        devices = loop.run_until_complete(discover())
-
-        for device in devices:
-            if addr:
-                if device.address == addr:
-                    self.device = device
-                    break
-            else:
-                try:
-                    if 1076 in device.metadata["manufacturer_data"].keys():
-                        self.device = device
-                        break
-                except KeyError:
-                    pass
+        scanner = BleakScanner() if scanner is None else scanner
+
+        if not address_or_ble_device:
+            device = loop.run_until_complete(
+                scanner.find_device_by_filter(
+                    lambda device, _: BT_MANUFACTURER_ID
+                    in device.metadata["manufacturer_data"].keys()
+                )
+            )
+        elif isinstance(address_or_ble_device, BLEDevice):
+            device = address_or_ble_device
         else:
+            device = loop.run_until_complete(
+                scanner.find_device_by_address(address_or_ble_device)
+            )
+
+        if device is None:
             raise RuntimeError(
-                "No address provided and could not find device via scan."
+                "No address or BLEDevice provided and could not find device via scan."
             )
 
+        self.device = device
+
         loop.run_until_complete(self._refresh_data())
 
     async def _send_command(self, command: str):
-        """ Send a command do the device.
+        """Send a command do the device.
 
         :param command: The command to send.
         """