Adjust entity registry access in integrations (1) (#88946)

homeassistant/components/rachio/entity.py

@@ -1,6 +1,6 @@
 """Adapter to wrap the rachiopy api for home assistant."""
 
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo, Entity
 
 from .const import DEFAULT_NAME, DOMAIN
@@ -25,7 +25,7 @@ class RachioDevice(Entity):
             },
             connections={
                 (
-                    device_registry.CONNECTION_NETWORK_MAC,
+                    dr.CONNECTION_NETWORK_MAC,
                     self._controller.mac_address,
                 )
             },

homeassistant/components/rainmachine/util.py

@@ -9,7 +9,7 @@ from typing import Any
 from homeassistant.backports.enum import StrEnum
 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,
     async_dispatcher_send,
@@ -55,7 +55,7 @@ def async_finish_entity_domain_replacements(
     entity_replacement_strategies: Iterable[EntityDomainReplacementStrategy],
 ) -> None:
     """Remove old entities and create a repairs issue with info on their replacement."""
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
     for strategy in entity_replacement_strategies:
         try:
             [registry_entry] = [

homeassistant/components/ruckus_unleashed/__init__.py

@@ -6,8 +6,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 
 from .const import (
     API_AP,
@@ -43,13 +42,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     system_info = await ruckus.system_info()
 
-    registry = device_registry.async_get(hass)
+    registry = dr.async_get(hass)
     ap_info = await ruckus.ap_info()
     for device in ap_info[API_AP][API_ID].values():
         registry.async_get_or_create(
             config_entry_id=entry.entry_id,
-            connections={(CONNECTION_NETWORK_MAC, device[API_MAC])},
+            connections={(dr.CONNECTION_NETWORK_MAC, device[API_MAC])},
-            identifiers={(CONNECTION_NETWORK_MAC, device[API_MAC])},
+            identifiers={(dr.CONNECTION_NETWORK_MAC, device[API_MAC])},
             manufacturer=MANUFACTURER,
             name=device[API_DEVICE_NAME],
             model=device[API_MODEL],

homeassistant/components/ruckus_unleashed/device_tracker.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 from homeassistant.components.device_tracker import ScannerEntity, SourceType
 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
 
@@ -37,7 +37,7 @@ async def async_setup_entry(
         coordinator.async_add_listener(router_update)
     )
 
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
     restore_entities(registry, coordinator, entry, async_add_entities, tracked)
 
 

homeassistant/components/sleepiq/entity.py

@@ -5,7 +5,7 @@ from typing import TypeVar
 from asyncsleepiq import SleepIQBed, SleepIQSleeper
 
 from homeassistant.core import callback
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo, Entity
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -21,7 +21,7 @@ _SleepIQCoordinatorT = TypeVar(
 def device_from_bed(bed: SleepIQBed) -> DeviceInfo:
     """Create a device given a bed."""
     return DeviceInfo(
-        connections={(device_registry.CONNECTION_NETWORK_MAC, bed.mac_addr)},
+        connections={(dr.CONNECTION_NETWORK_MAC, bed.mac_addr)},
         manufacturer="SleepNumber",
         name=bed.name,
         model=bed.model,

homeassistant/components/sonos/speaker.py

@@ -23,7 +23,7 @@ from homeassistant.components.media_player import DOMAIN as MP_DOMAIN
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import entity_registry as ent_reg
+from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
@@ -837,7 +837,7 @@ class SonosSpeaker:
                 # Skip updating existing single speakers in polling mode
                 return
 
-            entity_registry = ent_reg.async_get(self.hass)
+            entity_registry = er.async_get(self.hass)
             sonos_group = []
             sonos_group_entities = []
 

homeassistant/components/switcher_kis/__init__.py

@@ -12,7 +12,7 @@ from homeassistant.const import CONF_DEVICE_ID, EVENT_HOMEASSISTANT_STOP, Platfo
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.helpers import (
     config_validation as cv,
-    device_registry,
+    device_registry as dr,
     update_coordinator,
 )
 from homeassistant.helpers.dispatcher import async_dispatcher_send
@@ -165,10 +165,10 @@ class SwitcherDataUpdateCoordinator(
     @callback
     def async_setup(self) -> None:
         """Set up the coordinator."""
-        dev_reg = device_registry.async_get(self.hass)
+        dev_reg = dr.async_get(self.hass)
         dev_reg.async_get_or_create(
             config_entry_id=self.entry.entry_id,
-            connections={(device_registry.CONNECTION_NETWORK_MAC, self.mac_address)},
+            connections={(dr.CONNECTION_NETWORK_MAC, self.mac_address)},
             identifiers={(DOMAIN, self.device_id)},
             manufacturer="Switcher",
             name=self.name,

homeassistant/components/switcher_kis/button.py

@@ -19,7 +19,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
-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
@@ -132,9 +132,7 @@ class SwitcherThermostatButtonEntity(
         self._attr_name = f"{coordinator.name} {description.name}"
         self._attr_unique_id = f"{coordinator.mac_address}-{description.key}"
         self._attr_device_info = DeviceInfo(
-            connections={
+            connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-                (device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
-            }
         )
 
     async def async_press(self) -> None:

homeassistant/components/switcher_kis/climate.py

@@ -29,7 +29,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-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
@@ -94,9 +94,7 @@ class SwitcherClimateEntity(
         self._attr_name = coordinator.name
         self._attr_unique_id = f"{coordinator.device_id}-{coordinator.mac_address}"
         self._attr_device_info = DeviceInfo(
-            connections={
+            connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-                (device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
-            }
         )
 
         self._attr_min_temp = remote.min_temperature

homeassistant/components/switcher_kis/cover.py

@@ -17,7 +17,7 @@ from homeassistant.components.cover import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
-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
@@ -70,9 +70,7 @@ class SwitcherCoverEntity(
         self._attr_name = coordinator.name
         self._attr_unique_id = f"{coordinator.device_id}-{coordinator.mac_address}"
         self._attr_device_info = DeviceInfo(
-            connections={
+            connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-                (device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
-            }
         )
 
         self._update_data()

homeassistant/components/switcher_kis/sensor.py

@@ -13,7 +13,7 @@ from homeassistant.components.sensor import (
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfElectricCurrent, UnitOfPower
 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_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
@@ -118,9 +118,7 @@ class SwitcherSensorEntity(
             f"{coordinator.device_id}-{coordinator.mac_address}-{attribute}"
         )
         self._attr_device_info = {
-            "connections": {
+            "connections": {(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-                (device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
-            }
         }
 
     @property

homeassistant/components/switcher_kis/switch.py

@@ -15,7 +15,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import (
     config_validation as cv,
-    device_registry,
+    device_registry as dr,
     entity_platform,
 )
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -92,9 +92,7 @@ class SwitcherBaseSwitchEntity(
         self._attr_name = coordinator.name
         self._attr_unique_id = f"{coordinator.device_id}-{coordinator.mac_address}"
         self._attr_device_info = DeviceInfo(
-            connections={
+            connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
-                (device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
-            }
         )
 
     @callback

homeassistant/components/tplink_omada/entity.py

@@ -1,7 +1,7 @@
 """Base entity definitions."""
 from tplink_omada_client.devices import OmadaSwitch, OmadaSwitchPortDetails
 
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -25,7 +25,7 @@ class OmadaSwitchDeviceEntity(
     def device_info(self) -> DeviceInfo:
         """Return information about the device."""
         return DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, self.device.mac)},
+            connections={(dr.CONNECTION_NETWORK_MAC, self.device.mac)},
             identifiers={(DOMAIN, (self.device.mac))},
             manufacturer="TP-Link",
             model=self.device.model_display_name,

homeassistant/components/traccar/device_tracker.py

@@ -36,9 +36,8 @@ from homeassistant.const import (
     CONF_VERIFY_SSL,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import config_validation as cv, device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
-import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.event import async_track_time_interval
@@ -153,7 +152,7 @@ async def async_setup_entry(
     ] = async_dispatcher_connect(hass, TRACKER_UPDATE, _receive_data)
 
     # 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/tradfri/sensor.py

@@ -22,7 +22,7 @@ from homeassistant.const import (
     UnitOfTime,
 )
 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 .base_class import TradfriBaseEntity
@@ -108,7 +108,7 @@ SENSOR_DESCRIPTIONS_FAN: tuple[TradfriSensorEntityDescription, ...] = (
 @callback
 def _migrate_old_unique_ids(hass: HomeAssistant, old_unique_id: str, key: str) -> None:
     """Migrate unique IDs to the new format."""
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
 
     entity_id = ent_reg.async_get_entity_id(Platform.SENSOR, DOMAIN, old_unique_id)
 

homeassistant/components/upnp/__init__.py

@@ -12,7 +12,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 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 .const import (
     CONFIG_ENTRY_HOST,
@@ -118,11 +118,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     if device.serial_number:
         identifiers.add((IDENTIFIER_SERIAL_NUMBER, device.serial_number))
 
-    connections = {(device_registry.CONNECTION_UPNP, device.udn)}
+    connections = {(dr.CONNECTION_UPNP, device.udn)}
     if device_mac_address:
-        connections.add((device_registry.CONNECTION_NETWORK_MAC, device_mac_address))
+        connections.add((dr.CONNECTION_NETWORK_MAC, device_mac_address))
 
-    dev_registry = device_registry.async_get(hass)
+    dev_registry = dr.async_get(hass)
     device_entry = dev_registry.async_get_device(
         identifiers=identifiers, connections=connections
     )

homeassistant/components/velbus/__init__.py

@@ -13,9 +13,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ADDRESS, CONF_PORT, Platform
 from homeassistant.core import HomeAssistant, ServiceCall
 from homeassistant.exceptions import PlatformNotReady
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import config_validation as cv, device_registry as dr
-import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.device_registry import DeviceEntry
 from homeassistant.helpers.storage import STORAGE_DIR
 
 from .const import (
@@ -55,10 +53,8 @@ async def velbus_connect_task(
 
 def _migrate_device_identifiers(hass: HomeAssistant, entry_id: str) -> None:
     """Migrate old device indentifiers."""
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
-    devices: list[DeviceEntry] = device_registry.async_entries_for_config_entry(
+    devices: list[dr.DeviceEntry] = dr.async_entries_for_config_entry(dev_reg, entry_id)
-        dev_reg, entry_id
-    )
     for device in devices:
         old_identifier = list(next(iter(device.identifiers)))
         if len(old_identifier) > 2:

homeassistant/components/wiffi/__init__.py

@@ -9,7 +9,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PORT, CONF_TIMEOUT, Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
@@ -144,7 +144,7 @@ class WiffiEntity(Entity):
         """Initialize the base elements of a wiffi entity."""
         self._id = generate_unique_id(device, metric)
         self._device_info = DeviceInfo(
-            connections={(device_registry.CONNECTION_NETWORK_MAC, device.mac_address)},
+            connections={(dr.CONNECTION_NETWORK_MAC, device.mac_address)},
             identifiers={(DOMAIN, device.mac_address)},
             manufacturer="stall.biz",
             model=device.moduletype,

homeassistant/components/zwave_js/__init__.py

@@ -32,7 +32,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry, 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.dispatcher import async_dispatcher_send
 from homeassistant.helpers.issue_registry import (
@@ -161,8 +161,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     async_delete_issue(hass, DOMAIN, "invalid_server_version")
     LOGGER.info("Connected to Zwave JS Server")
 
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
     services = ZWaveServices(hass, ent_reg, dev_reg)
     services.async_register()
 
@@ -220,7 +220,7 @@ class DriverEvents:
     def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
         """Set up the driver events instance."""
         self.config_entry = entry
-        self.dev_reg = device_registry.async_get(hass)
+        self.dev_reg = dr.async_get(hass)
         self.hass = hass
         self.platform_setup_tasks: dict[str, asyncio.Task] = {}
         self.ready = asyncio.Event()
@@ -240,7 +240,7 @@ class DriverEvents:
             await driver.async_disable_statistics()
 
         # Check for nodes that no longer exist and remove them
-        stored_devices = device_registry.async_entries_for_config_entry(
+        stored_devices = dr.async_entries_for_config_entry(
             self.dev_reg, self.config_entry.entry_id
         )
         known_devices = [
@@ -311,7 +311,7 @@ class ControllerEvents:
         self.node_events = NodeEvents(hass, self)
 
     @callback
-    def remove_device(self, device: device_registry.DeviceEntry) -> None:
+    def remove_device(self, device: dr.DeviceEntry) -> None:
         """Remove device from registry."""
         # note: removal of entity registry entry is handled by core
         self.dev_reg.async_remove_device(device.id)
@@ -385,7 +385,7 @@ class ControllerEvents:
             self.remove_device(device)
 
     @callback
-    def register_node_in_dev_reg(self, node: ZwaveNode) -> device_registry.DeviceEntry:
+    def register_node_in_dev_reg(self, node: ZwaveNode) -> dr.DeviceEntry:
         """Register node in dev reg."""
         driver = self.driver_events.driver
         device_id = get_device_id(driver, node)
@@ -448,7 +448,7 @@ class NodeEvents:
         self.config_entry = controller_events.config_entry
         self.controller_events = controller_events
         self.dev_reg = controller_events.dev_reg
-        self.ent_reg = entity_registry.async_get(hass)
+        self.ent_reg = er.async_get(hass)
         self.hass = hass
 
     async def async_on_node_ready(self, node: ZwaveNode) -> None:
@@ -532,7 +532,7 @@ class NodeEvents:
 
     async def async_handle_discovery_info(
         self,
-        device: device_registry.DeviceEntry,
+        device: dr.DeviceEntry,
         disc_info: ZwaveDiscoveryInfo,
         value_updates_disc_info: dict[str, ZwaveDiscoveryInfo],
     ) -> None:

homeassistant/components/zwave_js/device_action.py

@@ -25,8 +25,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import Context, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
-from homeassistant.helpers import entity_registry
+from homeassistant.helpers import config_validation as cv, entity_registry as er
-import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.typing import ConfigType, TemplateVarsType
 
 from .config_validation import VALUE_SCHEMA
@@ -145,7 +144,7 @@ async def async_get_actions(
     hass: HomeAssistant, device_id: str
 ) -> list[dict[str, Any]]:
     """List device actions for Z-Wave JS devices."""
-    registry = entity_registry.async_get(hass)
+    registry = er.async_get(hass)
     actions: list[dict] = []
 
     node = async_get_node_from_device_id(hass, device_id)
@@ -179,7 +178,7 @@ async def async_get_actions(
 
     meter_endpoints: dict[int, dict[str, Any]] = defaultdict(dict)
 
-    for entry in entity_registry.async_entries_for_device(
+    for entry in er.async_entries_for_device(
         registry, device_id, include_disabled_entities=False
     ):
         # If an entry is unavailable, it is possible that the underlying value

homeassistant/components/zwave_js/device_trigger.py

@@ -22,8 +22,8 @@ from homeassistant.core import CALLBACK_TYPE, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import (
     config_validation as cv,
-    device_registry,
+    device_registry as dr,
-    entity_registry,
+    entity_registry as er,
 )
 from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 from homeassistant.helpers.typing import ConfigType
@@ -255,14 +255,14 @@ async def async_get_triggers(
         CONF_DOMAIN: DOMAIN,
     }
 
-    dev_reg = device_registry.async_get(hass)
+    dev_reg = dr.async_get(hass)
     node = async_get_node_from_device_id(hass, device_id, dev_reg)
 
     if node.client.driver and node.client.driver.controller.own_node == node:
         return triggers
 
     # We can add a node status trigger if the node status sensor is enabled
-    ent_reg = entity_registry.async_get(hass)
+    ent_reg = er.async_get(hass)
     entity_id = async_get_node_status_sensor_entity_id(
         hass, device_id, ent_reg, dev_reg
     )

homeassistant/components/zwave_me/__init__.py

@@ -7,8 +7,7 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_TOKEN, CONF_URL
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import device_registry
+from homeassistant.helpers import device_registry as dr
-from homeassistant.helpers.device_registry import DeviceRegistry
 from homeassistant.helpers.dispatcher import async_dispatcher_connect, dispatcher_send
 from homeassistant.helpers.entity import DeviceInfo, Entity
 
@@ -24,7 +23,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     controller = hass.data[DOMAIN][entry.entry_id] = ZWaveMeController(hass, entry)
     if await controller.async_establish_connection():
         await async_setup_platforms(hass, entry, controller)
-        registry = device_registry.async_get(hass)
+        registry = dr.async_get(hass)
         controller.remove_stale_devices(registry)
         return True
     raise ConfigEntryNotReady()
@@ -83,7 +82,7 @@ class ZWaveMeController:
         """Send signal to update device."""
         dispatcher_send(self._hass, f"ZWAVE_ME_INFO_{new_info.id}", new_info)
 
-    def remove_stale_devices(self, registry: DeviceRegistry):
+    def remove_stale_devices(self, registry: dr.DeviceRegistry):
         """Remove old-format devices in the registry."""
         for device_id in self.device_ids:
             device = registry.async_get_device(