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

This commit is contained in:
epenet 2023-03-01 08:02:34 +01:00 committed by GitHub
parent f2b736fad0
commit b3d6f098d2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
22 changed files with 61 additions and 79 deletions

View file

@ -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,
)
},

View file

@ -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] = [

View file

@ -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.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers import device_registry as dr
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])},
identifiers={(CONNECTION_NETWORK_MAC, device[API_MAC])},
connections={(dr.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],

View file

@ -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)

View file

@ -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,

View file

@ -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 = []

View file

@ -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,

View file

@ -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={
(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
}
connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
)
async def async_press(self) -> None:

View file

@ -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={
(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
}
connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
)
self._attr_min_temp = remote.min_temperature

View file

@ -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={
(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
}
connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
)
self._update_data()

View file

@ -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": {
(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
}
"connections": {(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
}
@property

View file

@ -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={
(device_registry.CONNECTION_NETWORK_MAC, coordinator.mac_address)
}
connections={(dr.CONNECTION_NETWORK_MAC, coordinator.mac_address)}
)
@callback

View file

@ -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,

View file

@ -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()

View file

@ -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)

View file

@ -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
)

View file

@ -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
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import DeviceEntry
from homeassistant.helpers import config_validation as cv, device_registry as dr
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)
devices: list[DeviceEntry] = device_registry.async_entries_for_config_entry(
dev_reg, entry_id
)
dev_reg = dr.async_get(hass)
devices: list[dr.DeviceEntry] = dr.async_entries_for_config_entry(dev_reg, entry_id)
for device in devices:
old_identifier = list(next(iter(device.identifiers)))
if len(old_identifier) > 2:

View file

@ -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,

View file

@ -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)
ent_reg = entity_registry.async_get(hass)
dev_reg = dr.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:

View file

@ -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
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import config_validation as cv, entity_registry as er
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

View file

@ -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,
entity_registry,
device_registry as dr,
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
)

View file

@ -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.device_registry import DeviceRegistry
from homeassistant.helpers import device_registry as dr
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(