From ad8436db2d5c6cfe49604293c2f6c1c7a83d087f Mon Sep 17 00:00:00 2001
From: Josh Wu <joshwu@google.com>
Date: Mon, 29 Jan 2024 16:42:45 +0800
Subject: [PATCH] Typing GATT Client

---
 bumble/gatt_client.py | 58 ++++++++++++++++++++++++++++---------------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/bumble/gatt_client.py b/bumble/gatt_client.py
index d5a8ec7c..1e764a45 100644
--- a/bumble/gatt_client.py
+++ b/bumble/gatt_client.py
@@ -136,7 +136,14 @@ def from_client(service_class, client: Client, service_uuid: UUID):
         service = services[0] if services else None
         return service_class(service) if service else None
 
-    def __init__(self, client, handle, end_group_handle, uuid, primary=True):
+    def __init__(
+        self,
+        client: Client,
+        handle: int,
+        end_group_handle: int,
+        uuid: UUID,
+        primary: bool = True,
+    ) -> None:
         attribute_type = (
             GATT_PRIMARY_SERVICE_ATTRIBUTE_TYPE
             if primary
@@ -146,10 +153,14 @@ def __init__(self, client, handle, end_group_handle, uuid, primary=True):
         self.uuid = uuid
         self.characteristics = []
 
-    async def discover_characteristics(self, uuids=()):
+    async def discover_characteristics(
+        self, uuids: Iterable[Union[str, UUID]] = ()
+    ) -> List[CharacteristicProxy]:
         return await self.client.discover_characteristics(uuids, self)
 
-    def get_characteristics_by_uuid(self, uuid):
+    def get_characteristics_by_uuid(
+        self, uuid: Union[str, UUID]
+    ) -> List[CharacteristicProxy]:
         return self.client.get_characteristics_by_uuid(uuid, self)
 
     def __str__(self) -> str:
@@ -163,10 +174,10 @@ class CharacteristicProxy(AttributeProxy):
 
     def __init__(
         self,
-        client,
-        handle,
-        end_group_handle,
-        uuid,
+        client: Client,
+        handle: int,
+        end_group_handle: int,
+        uuid: UUID,
         properties: int,
     ):
         super().__init__(client, handle, end_group_handle, uuid)
@@ -176,21 +187,26 @@ def __init__(
         self.descriptors_discovered = False
         self.subscribers = {}  # Map from subscriber to proxy subscriber
 
-    def get_descriptor(self, descriptor_type):
-        for descriptor in self.descriptors:
-            if descriptor.type == descriptor_type:
-                return descriptor
-
-        return None
+    def get_descriptor(
+        self, descriptor_type: Union[str, UUID]
+    ) -> Optional[DescriptorProxy]:
+        return next(
+            (
+                descriptor
+                for descriptor in self.descriptors
+                if descriptor.type == descriptor_type
+            ),
+            None,
+        )
 
-    async def discover_descriptors(self):
+    async def discover_descriptors(self) -> List[DescriptorProxy]:
         return await self.client.discover_descriptors(self)
 
     async def subscribe(
         self,
         subscriber: Optional[Callable[[bytes], Any]] = None,
         prefer_notify: bool = True,
-    ):
+    ) -> None:
         if subscriber is not None:
             if subscriber in self.subscribers:
                 # We already have a proxy subscriber
@@ -207,7 +223,9 @@ def on_change(value):
 
         return await self.client.subscribe(self, subscriber, prefer_notify)
 
-    async def unsubscribe(self, subscriber=None, force=False):
+    async def unsubscribe(
+        self, subscriber: Optional[Callable[[bytes], Any]] = None, force: bool = False
+    ) -> None:
         if subscriber in self.subscribers:
             subscriber = self.subscribers.pop(subscriber)
 
@@ -222,7 +240,7 @@ def __str__(self) -> str:
 
 
 class DescriptorProxy(AttributeProxy):
-    def __init__(self, client, handle, descriptor_type):
+    def __init__(self, client: Client, handle: int, descriptor_type: UUID) -> None:
         super().__init__(client, handle, 0, descriptor_type)
 
     def __str__(self) -> str:
@@ -339,11 +357,11 @@ async def request_mtu(self, mtu: int) -> int:
 
         return self.connection.att_mtu
 
-    def get_services_by_uuid(self, uuid: UUID) -> List[ServiceProxy]:
+    def get_services_by_uuid(self, uuid: Union[str, UUID]) -> List[ServiceProxy]:
         return [service for service in self.services if service.uuid == uuid]
 
     def get_characteristics_by_uuid(
-        self, uuid: UUID, service: Optional[ServiceProxy] = None
+        self, uuid: Union[str, UUID], service: Optional[ServiceProxy] = None
     ) -> List[CharacteristicProxy]:
         services = [service] if service else self.services
         return [
@@ -597,7 +615,7 @@ async def discover_included_services(
         return included_services
 
     async def discover_characteristics(
-        self, uuids, service: Optional[ServiceProxy]
+        self, uuids: Iterable[Union[str, UUID]], service: Optional[ServiceProxy]
     ) -> List[CharacteristicProxy]:
         '''
         See Vol 3, Part G - 4.6.1 Discover All Characteristics of a Service and 4.6.2