From 972eb34ed9ae526acba588286fab37dbbae47474 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Tue, 3 Jan 2023 08:19:53 +0100
Subject: [PATCH] Improve `bluetooth` generic typing (#84891)

Co-authored-by: J. Nick Koston <nick@koston.org>
---
 .../bluetooth/active_update_coordinator.py    |  2 +-
 .../bluetooth/passive_update_coordinator.py   | 22 ++++++---
 homeassistant/components/bsblan/climate.py    |  9 ++--
 homeassistant/components/elgato/light.py      |  2 +-
 .../components/keymitt_ble/entity.py          | 10 ++---
 .../components/keymitt_ble/switch.py          |  6 +--
 homeassistant/components/scrape/sensor.py     |  2 +-
 .../components/switchbot/coordinator.py       | 10 ++---
 homeassistant/components/switchbot/entity.py  |  5 ++-
 homeassistant/helpers/update_coordinator.py   | 45 ++++++++++++++-----
 10 files changed, 71 insertions(+), 42 deletions(-)

diff --git a/homeassistant/components/bluetooth/active_update_coordinator.py b/homeassistant/components/bluetooth/active_update_coordinator.py
index c4d40d5eaeb..5371d9f99fa 100644
--- a/homeassistant/components/bluetooth/active_update_coordinator.py
+++ b/homeassistant/components/bluetooth/active_update_coordinator.py
@@ -21,7 +21,7 @@ _T = TypeVar("_T")
 
 
 class ActiveBluetoothDataUpdateCoordinator(
-    Generic[_T], PassiveBluetoothDataUpdateCoordinator
+    PassiveBluetoothDataUpdateCoordinator, Generic[_T]
 ):
     """
     A coordinator that receives passive data from advertisements but can also poll.
diff --git a/homeassistant/components/bluetooth/passive_update_coordinator.py b/homeassistant/components/bluetooth/passive_update_coordinator.py
index 1eae49a6cab..6f1749aeef2 100644
--- a/homeassistant/components/bluetooth/passive_update_coordinator.py
+++ b/homeassistant/components/bluetooth/passive_update_coordinator.py
@@ -1,10 +1,13 @@
 """Passive update coordinator for the Bluetooth integration."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.helpers.update_coordinator import (
+    BaseCoordinatorEntity,
+    BaseDataUpdateCoordinatorProtocol,
+)
 
 from .update_coordinator import BasePassiveBluetoothCoordinator
 
@@ -14,8 +17,15 @@ if TYPE_CHECKING:
 
     from . import BluetoothChange, BluetoothScanningMode, BluetoothServiceInfoBleak
 
+_PassiveBluetoothDataUpdateCoordinatorT = TypeVar(
+    "_PassiveBluetoothDataUpdateCoordinatorT",
+    bound="PassiveBluetoothDataUpdateCoordinator",
+)
 
-class PassiveBluetoothDataUpdateCoordinator(BasePassiveBluetoothCoordinator):
+
+class PassiveBluetoothDataUpdateCoordinator(
+    BasePassiveBluetoothCoordinator, BaseDataUpdateCoordinatorProtocol
+):
     """Class to manage passive bluetooth advertisements.
 
     This coordinator is responsible for dispatching the bluetooth data
@@ -78,11 +88,11 @@ class PassiveBluetoothDataUpdateCoordinator(BasePassiveBluetoothCoordinator):
         self.async_update_listeners()
 
 
-class PassiveBluetoothCoordinatorEntity(CoordinatorEntity):
+class PassiveBluetoothCoordinatorEntity(
+    BaseCoordinatorEntity[_PassiveBluetoothDataUpdateCoordinatorT]
+):
     """A class for entities using DataUpdateCoordinator."""
 
-    coordinator: PassiveBluetoothDataUpdateCoordinator
-
     async def async_update(self) -> None:
         """All updates are passive."""
 
diff --git a/homeassistant/components/bsblan/climate.py b/homeassistant/components/bsblan/climate.py
index acf9ee25c57..cea58822ba8 100644
--- a/homeassistant/components/bsblan/climate.py
+++ b/homeassistant/components/bsblan/climate.py
@@ -64,10 +64,11 @@ async def async_setup_entry(
     )
 
 
-class BSBLANClimate(BSBLANEntity, CoordinatorEntity, ClimateEntity):
+class BSBLANClimate(
+    BSBLANEntity, CoordinatorEntity[DataUpdateCoordinator[State]], ClimateEntity
+):
     """Defines a BSBLAN climate device."""
 
-    coordinator: DataUpdateCoordinator[State]
     _attr_has_entity_name = True
     # Determine preset modes
     _attr_supported_features = (
@@ -80,7 +81,7 @@ class BSBLANClimate(BSBLANEntity, CoordinatorEntity, ClimateEntity):
 
     def __init__(
         self,
-        coordinator: DataUpdateCoordinator,
+        coordinator: DataUpdateCoordinator[State],
         client: BSBLAN,
         device: Device,
         info: Info,
@@ -89,7 +90,7 @@ class BSBLANClimate(BSBLANEntity, CoordinatorEntity, ClimateEntity):
     ) -> None:
         """Initialize BSBLAN climate device."""
         super().__init__(client, device, info, static, entry)
-        CoordinatorEntity.__init__(self, coordinator)
+        super(CoordinatorEntity, self).__init__(coordinator)
         self._attr_unique_id = f"{format_mac(device.MAC)}-climate"
 
         self._attr_min_temp = float(static.min_temp.value)
diff --git a/homeassistant/components/elgato/light.py b/homeassistant/components/elgato/light.py
index 2a9f63a83d7..6453950a814 100644
--- a/homeassistant/components/elgato/light.py
+++ b/homeassistant/components/elgato/light.py
@@ -76,7 +76,7 @@ class ElgatoLight(
     ) -> None:
         """Initialize Elgato Light."""
         super().__init__(client, info, mac)
-        CoordinatorEntity.__init__(self, coordinator)
+        super(CoordinatorEntity, self).__init__(coordinator)
 
         self._attr_min_mireds = 143
         self._attr_max_mireds = 344
diff --git a/homeassistant/components/keymitt_ble/entity.py b/homeassistant/components/keymitt_ble/entity.py
index dcda4a94027..31315e59efb 100644
--- a/homeassistant/components/keymitt_ble/entity.py
+++ b/homeassistant/components/keymitt_ble/entity.py
@@ -1,7 +1,7 @@
 """MicroBot class."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from homeassistant.components.bluetooth.passive_update_coordinator import (
     PassiveBluetoothCoordinatorEntity,
@@ -10,16 +10,12 @@ from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 
 from .const import MANUFACTURER
-
-if TYPE_CHECKING:
-    from . import MicroBotDataUpdateCoordinator
+from .coordinator import MicroBotDataUpdateCoordinator
 
 
-class MicroBotEntity(PassiveBluetoothCoordinatorEntity):
+class MicroBotEntity(PassiveBluetoothCoordinatorEntity[MicroBotDataUpdateCoordinator]):
     """Generic entity for all MicroBots."""
 
-    coordinator: MicroBotDataUpdateCoordinator
-
     def __init__(self, coordinator, config_entry):
         """Initialise the entity."""
         super().__init__(coordinator)
diff --git a/homeassistant/components/keymitt_ble/switch.py b/homeassistant/components/keymitt_ble/switch.py
index 92decea53ca..099ad1f228a 100644
--- a/homeassistant/components/keymitt_ble/switch.py
+++ b/homeassistant/components/keymitt_ble/switch.py
@@ -1,7 +1,7 @@
 """Switch platform for MicroBot."""
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import voluptuous as vol
 
@@ -11,11 +11,9 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, entity_platform
 
 from .const import DOMAIN
+from .coordinator import MicroBotDataUpdateCoordinator
 from .entity import MicroBotEntity
 
-if TYPE_CHECKING:
-    from . import MicroBotDataUpdateCoordinator
-
 CALIBRATE = "calibrate"
 CALIBRATE_SCHEMA = {
     vol.Required("depth"): cv.positive_int,
diff --git a/homeassistant/components/scrape/sensor.py b/homeassistant/components/scrape/sensor.py
index 22184a17b80..3d6d1db0ea3 100644
--- a/homeassistant/components/scrape/sensor.py
+++ b/homeassistant/components/scrape/sensor.py
@@ -203,7 +203,7 @@ class ScrapeSensor(CoordinatorEntity[ScrapeCoordinator], TemplateSensor):
         value_template: Template | None,
     ) -> None:
         """Initialize a web scrape sensor."""
-        CoordinatorEntity.__init__(self, coordinator)
+        super(CoordinatorEntity, self).__init__(coordinator)
         TemplateSensor.__init__(
             self,
             hass,
diff --git a/homeassistant/components/switchbot/coordinator.py b/homeassistant/components/switchbot/coordinator.py
index 77586c4202d..c12e8122e52 100644
--- a/homeassistant/components/switchbot/coordinator.py
+++ b/homeassistant/components/switchbot/coordinator.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 import asyncio
 import contextlib
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 import async_timeout
 import switchbot
@@ -25,9 +25,7 @@ _LOGGER = logging.getLogger(__name__)
 DEVICE_STARTUP_TIMEOUT = 30
 
 
-class SwitchbotDataUpdateCoordinator(
-    ActiveBluetoothDataUpdateCoordinator[dict[str, Any]]
-):
+class SwitchbotDataUpdateCoordinator(ActiveBluetoothDataUpdateCoordinator[None]):
     """Class to manage fetching switchbot data."""
 
     def __init__(
@@ -79,9 +77,9 @@ class SwitchbotDataUpdateCoordinator(
 
     async def _async_update(
         self, service_info: bluetooth.BluetoothServiceInfoBleak
-    ) -> dict[str, Any]:
+    ) -> None:
         """Poll the device."""
-        return await self.device.update()
+        await self.device.update()
 
     @callback
     def _async_handle_unavailable(
diff --git a/homeassistant/components/switchbot/entity.py b/homeassistant/components/switchbot/entity.py
index 60e7528dba6..c0e7a51170a 100644
--- a/homeassistant/components/switchbot/entity.py
+++ b/homeassistant/components/switchbot/entity.py
@@ -21,10 +21,11 @@ from .coordinator import SwitchbotDataUpdateCoordinator
 _LOGGER = logging.getLogger(__name__)
 
 
-class SwitchbotEntity(PassiveBluetoothCoordinatorEntity):
+class SwitchbotEntity(
+    PassiveBluetoothCoordinatorEntity[SwitchbotDataUpdateCoordinator]
+):
     """Generic entity encapsulating common features of Switchbot device."""
 
-    coordinator: SwitchbotDataUpdateCoordinator
     _device: SwitchbotDevice
     _attr_has_entity_name = True
 
diff --git a/homeassistant/helpers/update_coordinator.py b/homeassistant/helpers/update_coordinator.py
index 205a7848613..9c6747515a9 100644
--- a/homeassistant/helpers/update_coordinator.py
+++ b/homeassistant/helpers/update_coordinator.py
@@ -1,13 +1,14 @@
 """Helpers to help coordinate updates."""
 from __future__ import annotations
 
+from abc import abstractmethod
 import asyncio
 from collections.abc import Awaitable, Callable, Coroutine, Generator
 from datetime import datetime, timedelta
 import logging
 from random import randint
 from time import monotonic
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, Protocol, TypeVar
 import urllib.error
 
 import aiohttp
@@ -29,6 +30,9 @@ REQUEST_REFRESH_DEFAULT_COOLDOWN = 10
 REQUEST_REFRESH_DEFAULT_IMMEDIATE = True
 
 _T = TypeVar("_T")
+_BaseDataUpdateCoordinatorT = TypeVar(
+    "_BaseDataUpdateCoordinatorT", bound="BaseDataUpdateCoordinatorProtocol"
+)
 _DataUpdateCoordinatorT = TypeVar(
     "_DataUpdateCoordinatorT", bound="DataUpdateCoordinator[Any]"
 )
@@ -38,7 +42,17 @@ class UpdateFailed(Exception):
     """Raised when an update has failed."""
 
 
-class DataUpdateCoordinator(Generic[_T]):
+class BaseDataUpdateCoordinatorProtocol(Protocol):
+    """Base protocol type for DataUpdateCoordinator."""
+
+    @callback
+    def async_add_listener(
+        self, update_callback: CALLBACK_TYPE, context: Any = None
+    ) -> Callable[[], None]:
+        """Listen for data updates."""
+
+
+class DataUpdateCoordinator(BaseDataUpdateCoordinatorProtocol, Generic[_T]):
     """Class to manage fetching data from single endpoint."""
 
     def __init__(
@@ -346,11 +360,11 @@ class DataUpdateCoordinator(Generic[_T]):
         self.async_update_listeners()
 
 
-class CoordinatorEntity(entity.Entity, Generic[_DataUpdateCoordinatorT]):
-    """A class for entities using DataUpdateCoordinator."""
+class BaseCoordinatorEntity(entity.Entity, Generic[_BaseDataUpdateCoordinatorT]):
+    """Base class for all Coordinator entities."""
 
     def __init__(
-        self, coordinator: _DataUpdateCoordinatorT, context: Any = None
+        self, coordinator: _BaseDataUpdateCoordinatorT, context: Any = None
     ) -> None:
         """Create the entity with a DataUpdateCoordinator."""
         self.coordinator = coordinator
@@ -361,11 +375,6 @@ class CoordinatorEntity(entity.Entity, Generic[_DataUpdateCoordinatorT]):
         """No need to poll. Coordinator notifies entity of updates."""
         return False
 
-    @property
-    def available(self) -> bool:
-        """Return if entity is available."""
-        return self.coordinator.last_update_success
-
     async def async_added_to_hass(self) -> None:
         """When entity is added to hass."""
         await super().async_added_to_hass()
@@ -380,6 +389,22 @@ class CoordinatorEntity(entity.Entity, Generic[_DataUpdateCoordinatorT]):
         """Handle updated data from the coordinator."""
         self.async_write_ha_state()
 
+    @abstractmethod
+    async def async_update(self) -> None:
+        """Update the entity.
+
+        Only used by the generic entity update service.
+        """
+
+
+class CoordinatorEntity(BaseCoordinatorEntity[_DataUpdateCoordinatorT]):
+    """A class for entities using DataUpdateCoordinator."""
+
+    @property
+    def available(self) -> bool:
+        """Return if entity is available."""
+        return self.coordinator.last_update_success
+
     async def async_update(self) -> None:
         """Update the entity.