Adjust entity registry access in integrations (2) (#88947)

homeassistant/components/iotawatt/sensor.py

@@ -24,8 +24,7 @@ from homeassistant.const import (
     UnitOfPower,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import entity, entity_registry
+from homeassistant.helpers import device_registry as dr, entity, entity_registry as er
-from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -186,7 +185,9 @@ class IotaWattSensor(CoordinatorEntity[IotawattUpdater], SensorEntity):
     def device_info(self) -> entity.DeviceInfo:
         """Return device info."""
         return entity.DeviceInfo(
-            connections={(CONNECTION_NETWORK_MAC, self._sensor_data.hub_mac_address)},
+            connections={
+                (dr.CONNECTION_NETWORK_MAC, self._sensor_data.hub_mac_address)
+            },
             manufacturer="IoTaWatt",
             model="IoTaWatt",
         )
@@ -196,7 +197,7 @@ class IotaWattSensor(CoordinatorEntity[IotawattUpdater], SensorEntity):
         """Handle updated data from the coordinator."""
         if self._key not in self.coordinator.data["sensors"]:
             if self._attr_unique_id:
-                entity_registry.async_get(self.hass).async_remove(self.entity_id)
+                er.async_get(self.hass).async_remove(self.entity_id)
             else:
                 self.hass.async_create_task(self.async_remove())
             return

homeassistant/components/ipma/weather.py

@@ -43,7 +43,7 @@ from homeassistant.const import (
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.sun import is_up
 from homeassistant.util import Throttle
@@ -89,7 +89,7 @@ async def async_setup_entry(
 
     # Migrate old unique_id
     @callback
-    def _async_migrator(entity_entry: entity_registry.RegistryEntry):
+    def _async_migrator(entity_entry: er.RegistryEntry):
         # Reject if new unique_id
         if entity_entry.unique_id.count(",") == 2:
             return None
@@ -105,9 +105,7 @@ async def async_setup_entry(
         )
         return {"new_unique_id": new_unique_id}
 
-    await entity_registry.async_migrate_entries(
+    await er.async_migrate_entries(hass, config_entry.entry_id, _async_migrator)
-        hass, config_entry.entry_id, _async_migrator
-    )
 
     async_add_entities([IPMAWeather(location, api, config_entry.data)], True)
 

homeassistant/components/keenetic_ndms2/__init__.py

@@ -6,7 +6,7 @@ import logging
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_SCAN_INTERVAL, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry, entity_registry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from .const import (
     CONF_CONSIDER_HOME,
@@ -67,8 +67,8 @@ async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
         _LOGGER.debug(
             "Cleaning device_tracker entities since some interfaces are now untracked:"
         )
-        ent_reg = entity_registry.async_get(hass)
+        ent_reg = er.async_get(hass)
-        dev_reg = device_registry.async_get(hass)
+        dev_reg = dr.async_get(hass)
         # We keep devices currently connected to new_tracked_interfaces
         keep_devices: set[str] = {
             mac

homeassistant/components/keenetic_ndms2/device_tracker.py

@@ -12,7 +12,7 @@ from homeassistant.components.device_tracker import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
@@ -40,7 +40,7 @@ async def async_setup_entry(
 
     update_from_router()
 
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
     # Restore devices that are not a part of active clients list.
     restored = []
     for entity_entry in registry.entities.values():

homeassistant/components/kodi/device_trigger.py

@@ -13,7 +13,7 @@ from homeassistant.const import (
     CONF_TYPE,
 )
 from homeassistant.core import CALLBACK_TYPE, Event, HassJob, HomeAssistant, callback
-from homeassistant.helpers import config_validation as cv, entity_registry
+from homeassistant.helpers import config_validation as cv, entity_registry as er
 from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 from homeassistant.helpers.typing import ConfigType
 
@@ -33,11 +33,11 @@ async def async_get_triggers(
     hass: HomeAssistant, device_id: str
 ) -> list[dict[str, str]]:
     """List device triggers for Kodi devices."""
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
     triggers = []
 
     # Get all the integrations entities for this device
-    for entry in entity_registry.async_entries_for_device(registry, device_id):
+    for entry in er.async_entries_for_device(registry, device_id):
         if entry.domain == "media_player":
             triggers.append(
                 {

homeassistant/components/kodi/media_player.py

@@ -39,7 +39,7 @@ from homeassistant.const import (
 from homeassistant.core import CoreState, HomeAssistant, callback
 from homeassistant.helpers import (
     config_validation as cv,
-    device_registry,
+    device_registry as dr,
     entity_platform,
 )
 from homeassistant.helpers.entity import DeviceInfo
@@ -407,7 +407,7 @@ class KodiEntity(MediaPlayerEntity):
 
         version = (await self._kodi.get_application_properties(["version"]))["version"]
         sw_version = f"{version['major']}.{version['minor']}"
-        dev_reg = device_registry.async_get(self.hass)
+        dev_reg = dr.async_get(self.hass)
         device = dev_reg.async_get_device({(DOMAIN, self.unique_id)})
         dev_reg.async_update_device(device.id, sw_version=sw_version)
 

homeassistant/components/kraken/sensor.py

@@ -6,7 +6,7 @@ import logging
 from homeassistant.components.sensor import SensorEntity, SensorStateClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -55,11 +55,11 @@ async def async_setup_entry(
     @callback
     def async_update_sensors(hass: HomeAssistant, config_entry: ConfigEntry) -> None:
         """Add or remove sensors for configured tracked asset pairs."""
-        dev_reg = device_registry.async_get(hass)
+        dev_reg = dr.async_get(hass)
 
         existing_devices = {
             device.name: device.id
-            for device in device_registry.async_entries_for_config_entry(
+            for device in dr.async_entries_for_config_entry(
                 dev_reg, config_entry.entry_id
             )
         }
@@ -125,7 +125,7 @@ class KrakenSensor(
 
         self._attr_device_info = DeviceInfo(
             configuration_url="https://www.kraken.com/",
-            entry_type=device_registry.DeviceEntryType.SERVICE,
+            entry_type=dr.DeviceEntryType.SERVICE,
             identifiers={(DOMAIN, "_".join(self._device_name.split(" ")))},
             manufacturer="Kraken.com",
             name=self._device_name,

homeassistant/components/mazda/__init__.py

@@ -24,8 +24,11 @@ from homeassistant.exceptions import (
     ConfigEntryNotReady,
     HomeAssistantError,
 )
-from homeassistant.helpers import aiohttp_client, device_registry
+from homeassistant.helpers import (
-import homeassistant.helpers.config_validation as cv
+    aiohttp_client,
+    config_validation as cv,
+    device_registry as dr,
+)
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
@@ -81,7 +84,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     async def async_handle_service_call(service_call: ServiceCall) -> None:
         """Handle a service call."""
         # Get device entry from device registry
-        dev_reg = device_registry.async_get(hass)
+        dev_reg = dr.async_get(hass)
         device_id = service_call.data["device_id"]
         device_entry = dev_reg.async_get(device_id)
         if TYPE_CHECKING:
@@ -121,7 +124,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     def validate_mazda_device_id(device_id):
         """Check that a device ID exists in the registry and has at least one 'mazda' identifier."""
-        dev_reg = device_registry.async_get(hass)
+        dev_reg = dr.async_get(hass)
 
         if (device_entry := dev_reg.async_get(device_id)) is None:
             raise vol.Invalid("Invalid device ID")

homeassistant/components/metoffice/__init__.py

@@ -18,8 +18,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
-from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
@@ -53,7 +52,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     @callback
     def update_unique_id(
-        entity_entry: entity_registry.RegistryEntry,
+        entity_entry: er.RegistryEntry,
     ) -> dict[str, Any] | None:
         """Update unique ID of entity entry."""
 
@@ -86,7 +85,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
             }
         return None
 
-    await entity_registry.async_migrate_entries(hass, entry.entry_id, update_unique_id)
+    await er.async_migrate_entries(hass, entry.entry_id, update_unique_id)
 
     connection = datapoint.connection(api_key=api_key)
 
@@ -154,7 +153,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 def get_device_info(coordinates: str, name: str) -> DeviceInfo:
     """Return device registry information."""
     return DeviceInfo(
-        entry_type=DeviceEntryType.SERVICE,
+        entry_type=dr.DeviceEntryType.SERVICE,
         identifiers={(DOMAIN, coordinates)},
         manufacturer="Met Office",
         name=f"Met Office {name}",

homeassistant/components/mikrotik/device_tracker.py

@@ -10,7 +10,7 @@ from homeassistant.components.device_tracker import (
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 import homeassistant.util.dt as dt_util
@@ -31,7 +31,7 @@ async def async_setup_entry(
 
     tracked: dict[str, MikrotikDataUpdateCoordinatorTracker] = {}
 
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
 
     # Restore clients that is not a part of active clients list.
     for entity in registry.entities.values():

homeassistant/components/nam/__init__.py

@@ -21,9 +21,8 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
@@ -70,7 +69,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     # Remove air_quality entities from registry if they exist
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
     for sensor_type in ("sds", ATTR_SDS011, ATTR_SPS30):
         unique_id = f"{coordinator.unique_id}-{sensor_type}"
         if entity_id := ent_reg.async_get_entity_id(
@@ -130,7 +129,7 @@ class NAMDataUpdateCoordinator(DataUpdateCoordinator[NAMSensors]):
     def device_info(self) -> DeviceInfo:
         """Return the device info."""
         return DeviceInfo(
-            connections={(CONNECTION_NETWORK_MAC, cast(str, self._unique_id))},
+            connections={(dr.CONNECTION_NETWORK_MAC, cast(str, self._unique_id))},
             name="Nettigo Air Monitor",
             sw_version=self.nam.software_version,
             manufacturer=MANUFACTURER,

homeassistant/components/nam/sensor.py

@@ -26,7 +26,7 @@ from homeassistant.const import (
     UnitOfTemperature,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -369,7 +369,7 @@ async def async_setup_entry(
 
     # Due to the change of the attribute name of two sensors, it is necessary to migrate
     # the unique_ids to the new names.
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
     for old_sensor, new_sensor in MIGRATION_SENSORS:
         old_unique_id = f"{coordinator.unique_id}-{old_sensor}"
         new_unique_id = f"{coordinator.unique_id}-{new_sensor}"

homeassistant/components/netatmo/device_trigger.py

@@ -20,7 +20,7 @@ from homeassistant.core import CALLBACK_TYPE, HomeAssistant
 from homeassistant.helpers import (
     config_validation as cv,
     device_registry as dr,
-    entity_registry,
+    entity_registry as er,
 )
 from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 from homeassistant.helpers.typing import ConfigType
@@ -93,11 +93,11 @@ async def async_get_triggers(
     hass: HomeAssistant, device_id: str
 ) -> list[dict[str, str]]:
     """List device triggers for Netatmo devices."""
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
     device_registry = dr.async_get(hass)
     triggers = []
 
-    for entry in entity_registry.async_entries_for_device(registry, device_id):
+    for entry in er.async_entries_for_device(registry, device_id):
         if (
             device := device_registry.async_get(device_id)
         ) is None or device.model is None:

homeassistant/components/nobo_hub/__init__.py

@@ -11,7 +11,7 @@ from homeassistant.const import (
     Platform,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 
 from .const import (
     ATTR_HARDWARE_VERSION,
@@ -38,7 +38,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     hass.data.setdefault(DOMAIN, {})
 
     # Register hub as device
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
     dev_reg.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={(DOMAIN, hub.hub_info[ATTR_SERIAL])},

homeassistant/components/nuki/__init__.py

@@ -16,7 +16,7 @@ from homeassistant import exceptions
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PORT, CONF_TOKEN, Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import device_registry, entity_registry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.update_coordinator import (
     CoordinatorEntity,
     DataUpdateCoordinator,
@@ -78,7 +78,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # Device registration for the bridge
     info = bridge.info()
     bridge_id = parse_id(info["ids"]["hardwareId"])
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
     dev_reg.async_get_or_create(
         config_entry_id=entry.entry_id,
         identifiers={(DOMAIN, bridge_id)},
@@ -150,7 +150,7 @@ class NukiCoordinator(DataUpdateCoordinator[None]):
         except RequestException as err:
             raise UpdateFailed(f"Error communicating with Bridge: {err}") from err
 
-        ent_reg = entity_registry.async_get(self.hass)
+        ent_reg = er.async_get(self.hass)
         for event, device_ids in events.items():
             for device_id in device_ids:
                 entity_id = ent_reg.async_get_entity_id(

homeassistant/components/owntracks/device_tracker.py

@@ -13,7 +13,7 @@ from homeassistant.const import (
     ATTR_LONGITUDE,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
@@ -26,7 +26,7 @@ async def async_setup_entry(
 ) -> None:
     """Set up OwnTracks based off an entry."""
     # Restore previously loaded devices
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
     dev_ids = {
         identifier[1]
         for device in dev_reg.devices.values()

homeassistant/components/plex/__init__.py

@@ -18,7 +18,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_URL, CONF_VERIFY_SSL, EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers import device_registry as dev_reg, entity_registry as ent_reg
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.debounce import Debouncer
 from homeassistant.helpers.dispatcher import (
@@ -289,17 +289,15 @@ async def async_options_updated(hass: HomeAssistant, entry: ConfigEntry) -> None
 @callback
 def async_cleanup_plex_devices(hass, entry):
     """Clean up old and invalid devices from the registry."""
-    device_registry = dev_reg.async_get(hass)
+    device_registry = dr.async_get(hass)
-    entity_registry = ent_reg.async_get(hass)
+    entity_registry = er.async_get(hass)
 
-    device_entries = dev_reg.async_entries_for_config_entry(
+    device_entries = dr.async_entries_for_config_entry(device_registry, entry.entry_id)
-        device_registry, entry.entry_id
-    )
 
     for device_entry in device_entries:
         if (
             len(
-                ent_reg.async_entries_for_device(
+                er.async_entries_for_device(
                     entity_registry, device_entry.id, include_disabled_entities=True
                 )
             )

homeassistant/components/point/__init__.py

@@ -18,7 +18,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import config_validation as cv, device_registry
+from homeassistant.helpers import config_validation as cv, device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
@@ -322,9 +322,7 @@ class MinutPointEntity(Entity):
         """Return a device description for device registry."""
         device = self.device.device
         return DeviceInfo(
-            connections={
+            connections={(dr.CONNECTION_NETWORK_MAC, device["device_mac"])},
-                (device_registry.CONNECTION_NETWORK_MAC, device["device_mac"])
-            },
             identifiers={(DOMAIN, device["device_id"])},
             manufacturer="Minut",
             model=f"Point v{device['hardware_version']}",

homeassistant/components/ps4/__init__.py

@@ -23,7 +23,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, ServiceCall, split_entity_id
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import config_validation as cv, entity_registry
+from homeassistant.helpers import config_validation as cv, entity_registry as er
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.json import save_json
 from homeassistant.helpers.typing import ConfigType
@@ -116,7 +116,7 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     # Migrate Version 2 -> Version 3: Update identifier format.
     if version == 2:
         # Prevent changing entity_id. Updates entity registry.
-        registry = entity_registry.async_get(hass)
+        registry = er.async_get(hass)
 
         for entity_id, e_entry in registry.entities.items():
             if e_entry.config_entry_id == entry.entry_id:

homeassistant/components/ps4/media_player.py

@@ -24,7 +24,7 @@ from homeassistant.const import (
     CONF_TOKEN,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry, entity_registry
+from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -321,8 +321,8 @@ class PS4Device(MediaPlayerEntity):
         # If cannot get status on startup, assume info from registry.
         if status is None:
             _LOGGER.info("Assuming status from registry")
-            e_registry = entity_registry.async_get(self.hass)
+            e_registry = er.async_get(self.hass)
-            d_registry = device_registry.async_get(self.hass)
+            d_registry = dr.async_get(self.hass)
             for entity_id, entry in e_registry.entities.items():
                 if entry.config_entry_id == self._entry_id:
                     self._attr_unique_id = entry.unique_id