Shelly entity device info code quality (home-assistant#149477)

Author: thecode

homeassistant/components/shelly/button.py

@@ -25,13 +25,8 @@
 
 from .const import DOMAIN, LOGGER, SHELLY_GAS_MODELS
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
-from .utils import (
-    get_block_device_info,
-    get_blu_trv_device_info,
-    get_device_entry_gen,
-    get_rpc_device_info,
-    get_rpc_key_ids,
-)
+from .entity import get_entity_block_device_info, get_entity_rpc_device_info
+from .utils import get_blu_trv_device_info, get_device_entry_gen, get_rpc_key_ids
 
 PARALLEL_UPDATES = 0
 
@@ -233,23 +228,9 @@ def __init__(
 
         self._attr_unique_id = f"{coordinator.mac}_{description.key}"
         if isinstance(coordinator, ShellyBlockCoordinator):
-            self._attr_device_info = get_block_device_info(
-                coordinator.device,
-                coordinator.mac,
-                coordinator.configuration_url,
-                coordinator.model,
-                coordinator.model_name,
-                suggested_area=coordinator.suggested_area,
-            )
+            self._attr_device_info = get_entity_block_device_info(coordinator)
         else:
-            self._attr_device_info = get_rpc_device_info(
-                coordinator.device,
-                coordinator.mac,
-                coordinator.configuration_url,
-                coordinator.model,
-                coordinator.model_name,
-                suggested_area=coordinator.suggested_area,
-            )
+            self._attr_device_info = get_entity_rpc_device_info(coordinator)
 
     async def _press_method(self) -> None:
         """Press method."""

homeassistant/components/shelly/climate.py

@@ -38,10 +38,9 @@
     SHTRV_01_TEMPERATURE_SETTINGS,
 )
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
-from .entity import ShellyRpcEntity, rpc_call
+from .entity import ShellyRpcEntity, get_entity_block_device_info, rpc_call
 from .utils import (
     async_remove_shelly_entity,
-    get_block_device_info,
     get_block_entity_name,
     get_blu_trv_device_info,
     get_device_entry_gen,
@@ -210,15 +209,7 @@ def __init__(
             ]
         elif entry is not None:
             self._unique_id = entry.unique_id
-        self._attr_device_info = get_block_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            sensor_block,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_block_device_info(coordinator, sensor_block)
         self._attr_name = get_block_entity_name(
             self.coordinator.device, sensor_block, None
         )

homeassistant/components/shelly/entity.py

@@ -13,6 +13,7 @@
 from homeassistant.core import HomeAssistant, State, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import entity_registry as er
+from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity, EntityDescription
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.entity_registry import RegistryEntry
@@ -368,15 +369,7 @@ def __init__(self, coordinator: ShellyBlockCoordinator, block: Block) -> None:
         super().__init__(coordinator)
         self.block = block
         self._attr_name = get_block_entity_name(coordinator.device, block)
-        self._attr_device_info = get_block_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            block,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_block_device_info(coordinator, block)
         self._attr_unique_id = f"{coordinator.mac}-{block.description}"
 
     # pylint: disable-next=hass-missing-super-call
@@ -417,15 +410,7 @@ def __init__(self, coordinator: ShellyRpcCoordinator, key: str) -> None:
         """Initialize Shelly entity."""
         super().__init__(coordinator)
         self.key = key
-        self._attr_device_info = get_rpc_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            key,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_rpc_device_info(coordinator, key)
         self._attr_unique_id = f"{coordinator.mac}-{key}"
         self._attr_name = get_rpc_entity_name(coordinator.device, key)
 
@@ -539,14 +524,7 @@ def __init__(
             coordinator.device, None, description.name
         )
         self._attr_unique_id = f"{coordinator.mac}-{attribute}"
-        self._attr_device_info = get_block_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_block_device_info(coordinator)
         self._last_value = None
 
     @property
@@ -653,15 +631,7 @@ def __init__(
         self.block: Block | None = block  # type: ignore[assignment]
         self.entity_description = description
 
-        self._attr_device_info = get_block_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            block,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_block_device_info(coordinator, block)
 
         if block is not None:
             self._attr_unique_id = (
@@ -726,18 +696,8 @@ def __init__(
         self.attribute = attribute
         self.entity_description = description
 
-        self._attr_device_info = get_rpc_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            key,
-            suggested_area=coordinator.suggested_area,
-        )
-        self._attr_unique_id = self._attr_unique_id = (
-            f"{coordinator.mac}-{key}-{attribute}"
-        )
+        self._attr_device_info = get_entity_rpc_device_info(coordinator, key)
+        self._attr_unique_id = f"{coordinator.mac}-{key}-{attribute}"
         self._last_value = None
 
         if coordinator.device.initialized:
@@ -763,3 +723,37 @@ def get_entity_class(
         return description.entity_class
 
     return sensor_class
+
+
+def get_entity_block_device_info(
+    coordinator: ShellyBlockCoordinator,
+    block: Block | None = None,
+) -> DeviceInfo:
+    """Get device info for block entities."""
+    return get_block_device_info(
+        coordinator.device,
+        coordinator.mac,
+        coordinator.configuration_url,
+        coordinator.model,
+        coordinator.model_name,
+        block,
+        suggested_area=coordinator.suggested_area,
+    )
+
+
+def get_entity_rpc_device_info(
+    coordinator: ShellyRpcCoordinator,
+    key: str | None = None,
+    emeter_phase: str | None = None,
+) -> DeviceInfo:
+    """Get device info for RPC entities."""
+    return get_rpc_device_info(
+        coordinator.device,
+        coordinator.mac,
+        coordinator.configuration_url,
+        coordinator.model,
+        coordinator.model_name,
+        key,
+        emeter_phase=emeter_phase,
+        suggested_area=coordinator.suggested_area,
+    )

homeassistant/components/shelly/event.py

@@ -26,12 +26,11 @@
     SHIX3_1_INPUTS_EVENTS_TYPES,
 )
 from .coordinator import ShellyBlockCoordinator, ShellyConfigEntry, ShellyRpcCoordinator
-from .entity import ShellyBlockEntity
+from .entity import ShellyBlockEntity, get_entity_rpc_device_info
 from .utils import (
     async_remove_orphaned_entities,
     async_remove_shelly_entity,
     get_device_entry_gen,
-    get_rpc_device_info,
     get_rpc_entity_name,
     get_rpc_key_instances,
     is_block_momentary_input,
@@ -206,15 +205,7 @@ def __init__(
         """Initialize Shelly entity."""
         super().__init__(coordinator)
         self.event_id = int(key.split(":")[-1])
-        self._attr_device_info = get_rpc_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            key,
-            suggested_area=coordinator.suggested_area,
-        )
+        self._attr_device_info = get_entity_rpc_device_info(coordinator, key)
         self._attr_unique_id = f"{coordinator.mac}-{key}"
         self._attr_name = get_rpc_entity_name(coordinator.device, key)
         self.entity_description = description

homeassistant/components/shelly/sensor.py

@@ -52,13 +52,13 @@
     async_setup_entry_attribute_entities,
     async_setup_entry_rest,
     async_setup_entry_rpc,
+    get_entity_rpc_device_info,
 )
 from .utils import (
     async_remove_orphaned_entities,
     get_blu_trv_device_info,
     get_device_entry_gen,
     get_device_uptime,
-    get_rpc_device_info,
     get_shelly_air_lamp_life,
     get_virtual_component_ids,
     is_rpc_wifi_stations_disabled,
@@ -138,15 +138,8 @@ def __init__(
         """Initialize select."""
         super().__init__(coordinator, key, attribute, description)
 
-        self._attr_device_info = get_rpc_device_info(
-            coordinator.device,
-            coordinator.mac,
-            coordinator.configuration_url,
-            coordinator.model,
-            coordinator.model_name,
-            key,
-            emeter_phase=description.emeter_phase,
-            suggested_area=coordinator.suggested_area,
+        self._attr_device_info = get_entity_rpc_device_info(
+            coordinator, key, emeter_phase=description.emeter_phase
         )