Add typing to IIDManager (#468)

pyhap/iid_manager.py

@@ -1,19 +1,26 @@
 """Module for the IIDManager class."""
 import logging
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+if TYPE_CHECKING:
+    from .characteristic import Characteristic
+    from .service import Service
+
+    ServiceOrCharType = Union[Service, Characteristic]
 
 logger = logging.getLogger(__name__)
 
 
 class IIDManager:
     """Maintains a mapping between Service/Characteristic objects and IIDs."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize an empty instance."""
         self.counter = 0
-        self.iids = {}
-        self.objs = {}
+        self.iids: Dict["ServiceOrCharType", int] = {}
+        self.objs: Dict[int, "ServiceOrCharType"] = {}
 
-    def assign(self, obj):
+    def assign(self, obj: "ServiceOrCharType") -> None:
         """Assign an IID to given object. Print warning if already assigned.
 
         :param obj: The object that will be assigned an IID.
@@ -32,23 +39,23 @@ def assign(self, obj):
         self.iids[obj] = iid
         self.objs[iid] = obj
 
-    def get_iid_for_obj(self, obj):
+    def get_iid_for_obj(self, obj: "ServiceOrCharType") -> int:
         """Get the IID for the given object.
 
         Override this method to provide custom IID assignment.
         """
         self.counter += 1
         return self.counter
 
-    def get_obj(self, iid):
+    def get_obj(self, iid: int) -> "ServiceOrCharType":
         """Get the object that is assigned the given IID."""
         return self.objs.get(iid)
 
-    def get_iid(self, obj):
+    def get_iid(self, obj: "ServiceOrCharType") -> int:
         """Get the IID assigned to the given object."""
         return self.iids.get(obj)
 
-    def remove_obj(self, obj):
+    def remove_obj(self, obj: "ServiceOrCharType") -> Optional[int]:
         """Remove an object from the IID list."""
         iid = self.iids.pop(obj, None)
         if iid is None:
@@ -57,7 +64,7 @@ def remove_obj(self, obj):
         del self.objs[iid]
         return iid
 
-    def remove_iid(self, iid):
+    def remove_iid(self, iid: int) -> Optional["ServiceOrCharType"]:
         """Remove an object with an IID from the IID list."""
         obj = self.objs.pop(iid, None)
         if obj is None: