Elgato refactoring; add base entity

homeassistant/components/elgato/__init__.py

@@ -1,7 +1,8 @@
 """Support for Elgato Lights."""
 import logging
+from typing import NamedTuple
 
-from elgato import Elgato, ElgatoConnectionError
+from elgato import Elgato, ElgatoConnectionError, Info
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT, Platform
@@ -14,6 +15,13 @@
 PLATFORMS = [Platform.BUTTON, Platform.LIGHT]
 
 
+class HomeAssistantElgatoData(NamedTuple):
+    """Elgato data stored in the Home Assistant data object."""
+
+    client: Elgato
+    info: Info
+
+
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Elgato Light from a config entry."""
     session = async_get_clientsession(hass)
@@ -25,22 +33,24 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     # Ensure we can connect to it
     try:
-        await elgato.info()
+        info = await elgato.info()
     except ElgatoConnectionError as exception:
         logging.getLogger(__name__).debug("Unable to connect: %s", exception)
         raise ConfigEntryNotReady from exception
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = elgato
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = HomeAssistantElgatoData(
+        client=elgato,
+        info=info,
+    )
+
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
 
 
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload Elgato Light config entry."""
-    # Unload entities for this entry/device.
-    unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-    if unload_ok:
+    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         # Cleanup
         del hass.data[DOMAIN][entry.entry_id]
         if not hass.data[DOMAIN]:

homeassistant/components/elgato/button.py

@@ -8,10 +8,12 @@
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import DeviceInfo, EntityCategory
+from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import HomeAssistantElgatoData
 from .const import DOMAIN
+from .entity import ElgatoEntity
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -22,18 +24,16 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Elgato button based on a config entry."""
-    elgato: Elgato = hass.data[DOMAIN][entry.entry_id]
-    info = await elgato.info()
-    async_add_entities([ElgatoIdentifyButton(elgato, info)])
+    data: HomeAssistantElgatoData = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([ElgatoIdentifyButton(data.client, data.info)])
 
 
-class ElgatoIdentifyButton(ButtonEntity):
+class ElgatoIdentifyButton(ElgatoEntity, ButtonEntity):
     """Defines an Elgato identify button."""
 
-    def __init__(self, elgato: Elgato, info: Info) -> None:
+    def __init__(self, client: Elgato, info: Info) -> None:
         """Initialize the button entity."""
-        self.elgato = elgato
-        self._info = info
+        super().__init__(client, info)
         self.entity_description = ButtonEntityDescription(
             key="identify",
             name="Identify",
@@ -42,20 +42,9 @@ def __init__(self, elgato: Elgato, info: Info) -> None:
         )
         self._attr_unique_id = f"{info.serial_number}_{self.entity_description.key}"
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return device information about this Elgato Light."""
-        return DeviceInfo(
-            identifiers={(DOMAIN, self._info.serial_number)},
-            manufacturer="Elgato",
-            model=self._info.product_name,
-            name=self._info.product_name,
-            sw_version=f"{self._info.firmware_version} ({self._info.firmware_build_number})",
-        )
-
     async def async_press(self) -> None:
         """Identify the light, will make it blink."""
         try:
-            await self.elgato.identify()
+            await self.client.identify()
         except ElgatoError:
             _LOGGER.exception("An error occurred while identifying the Elgato Light")

homeassistant/components/elgato/entity.py

@@ -0,0 +1,22 @@
+"""Base entity for the Elgato integration."""
+
+from elgato import Elgato, Info
+
+from homeassistant.helpers.entity import DeviceInfo, Entity
+
+from .const import DOMAIN
+
+
+class ElgatoEntity(Entity):
+    """Defines an Elgato entity."""
+
+    def __init__(self, client: Elgato, info: Info) -> None:
+        """Initialize an Elgato entity."""
+        self.client = client
+        self._attr_device_info = DeviceInfo(
+            identifiers={(DOMAIN, info.serial_number)},
+            manufacturer="Elgato",
+            model=info.product_name,
+            name=info.product_name,
+            sw_version=f"{info.firmware_version} ({info.firmware_build_number})",
+        )

homeassistant/components/elgato/light.py

@@ -17,13 +17,14 @@
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import (
     AddEntitiesCallback,
     async_get_current_platform,
 )
 
+from . import HomeAssistantElgatoData
 from .const import DOMAIN, SERVICE_IDENTIFY
+from .entity import ElgatoEntity
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -37,10 +38,9 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Elgato Light based on a config entry."""
-    elgato: Elgato = hass.data[DOMAIN][entry.entry_id]
-    info = await elgato.info()
-    settings = await elgato.settings()
-    async_add_entities([ElgatoLight(elgato, info, settings)], True)
+    data: HomeAssistantElgatoData = hass.data[DOMAIN][entry.entry_id]
+    settings = await data.client.settings()
+    async_add_entities([ElgatoLight(data.client, data.info, settings)], True)
 
     platform = async_get_current_platform()
     platform.async_register_entity_service(
@@ -50,15 +50,13 @@ async def async_setup_entry(
     )
 
 
-class ElgatoLight(LightEntity):
+class ElgatoLight(ElgatoEntity, LightEntity):
     """Defines an Elgato Light."""
 
-    def __init__(self, elgato: Elgato, info: Info, settings: Settings) -> None:
+    def __init__(self, client: Elgato, info: Info, settings: Settings) -> None:
         """Initialize Elgato Light."""
-        self._info = info
-        self._settings = settings
+        super().__init__(client, info)
         self._state: State | None = None
-        self.elgato = elgato
 
         min_mired = 143
         max_mired = 344
@@ -122,7 +120,7 @@ def is_on(self) -> bool:
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the light."""
         try:
-            await self.elgato.light(on=False)
+            await self.client.light(on=False)
         except ElgatoError:
             _LOGGER.error("An error occurred while updating the Elgato Light")
             self._state = None
@@ -154,7 +152,7 @@ async def async_turn_on(self, **kwargs: Any) -> None:
             temperature = self.color_temp
 
         try:
-            await self.elgato.light(
+            await self.client.light(
                 on=True,
                 brightness=brightness,
                 hue=hue,
@@ -169,29 +167,18 @@ async def async_update(self) -> None:
         """Update Elgato entity."""
         restoring = self._state is None
         try:
-            self._state = await self.elgato.state()
+            self._state = await self.client.state()
             if restoring:
                 _LOGGER.info("Connection restored")
         except ElgatoError as err:
             meth = _LOGGER.error if self._state else _LOGGER.debug
             meth("An error occurred while updating the Elgato Light: %s", err)
             self._state = None
 
-    @property
-    def device_info(self) -> DeviceInfo:
-        """Return device information about this Elgato Light."""
-        return DeviceInfo(
-            identifiers={(DOMAIN, self._info.serial_number)},
-            manufacturer="Elgato",
-            model=self._info.product_name,
-            name=self._info.product_name,
-            sw_version=f"{self._info.firmware_version} ({self._info.firmware_build_number})",
-        )
-
     async def async_identify(self) -> None:
         """Identify the light, will make it blink."""
         try:
-            await self.elgato.identify()
+            await self.client.identify()
         except ElgatoError:
             _LOGGER.exception("An error occurred while identifying the Elgato Light")
             self._state = None