diff --git a/pyhap/iid_manager.py b/pyhap/iid_manager.py index 20ce92c6..9bee16f8 100644 --- a/pyhap/iid_manager.py +++ b/pyhap/iid_manager.py @@ -9,6 +9,8 @@ from .characteristic import Characteristic from .service import Service + ServiceOrCharType = Union[Service, Characteristic] + class IIDManager: """Maintains a mapping between Service/Characteristic objects and IIDs.""" @@ -16,10 +18,10 @@ class IIDManager: def __init__(self) -> None: """Initialize an empty instance.""" self.counter = 0 - self.iids: Dict[Union[Service, Characteristic], int] = {} - self.objs: Dict[int, Union[Service, Characteristic]] = {} + self.iids: Dict["ServiceOrCharType", int] = {} + self.objs: Dict[int, "ServiceOrCharType"] = {} - def assign(self, obj: Union[Service, Characteristic]) -> None: + 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. @@ -38,7 +40,7 @@ def assign(self, obj: Union[Service, Characteristic]) -> None: self.iids[obj] = iid self.objs[iid] = obj - def get_iid_for_obj(self, obj: Union[Service, Characteristic]) -> int: + def get_iid_for_obj(self, obj: "ServiceOrCharType") -> int: """Get the IID for the given object. Override this method to provide custom IID assignment. @@ -46,15 +48,15 @@ def get_iid_for_obj(self, obj: Union[Service, Characteristic]) -> int: self.counter += 1 return self.counter - def get_obj(self, iid: int) -> Union[Service, Characteristic]: + 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: Union[Service, Characteristic]) -> int: + 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: Union[Service, Characteristic]) -> Optional[int]: + 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: @@ -63,7 +65,7 @@ def remove_obj(self, obj: Union[Service, Characteristic]) -> Optional[int]: del self.objs[iid] return iid - def remove_iid(self, iid: int) -> Optional[Union[Service, Characteristic]]: + 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: