Adjust entity registry access in integrations (3) (#88948)

This commit is contained in:
epenet 2023-03-01 08:24:56 +01:00 committed by GitHub
parent f69d76702a
commit 42b74e7f56
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
21 changed files with 90 additions and 93 deletions

View file

@ -20,13 +20,12 @@ from homeassistant.components.camera import (
from homeassistant.components.ffmpeg import FFmpegManager, get_ffmpeg_manager
from homeassistant.const import ATTR_ENTITY_ID, CONF_NAME, STATE_OFF, STATE_ON
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry
from homeassistant.helpers import config_validation as cv, entity_registry as er
from homeassistant.helpers.aiohttp_client import (
async_aiohttp_proxy_stream,
async_aiohttp_proxy_web,
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.typing import ConfigType, DiscoveryInfoType
@ -146,7 +145,7 @@ async def async_setup_platform(
# with this version, update the old entity with the new unique id.
serial_number = await device.api.async_serial_number
serial_number = serial_number.strip()
registry = entity_registry.async_get(hass)
registry = er.async_get(hass)
entity_id = registry.async_get_entity_id(CAMERA_DOMAIN, DOMAIN, serial_number)
if entity_id is not None:
_LOGGER.debug("Updating unique id for camera %s", entity_id)

View file

@ -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
@ -32,11 +32,11 @@ async def async_get_triggers(
hass: HomeAssistant, device_id: str
) -> list[dict[str, str]]:
"""List device triggers for Arcam FMJ Receiver control 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(
{

View file

@ -10,8 +10,7 @@ from coinbase.wallet.error import AuthenticationError
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_API_KEY, CONF_API_TOKEN, Platform
from homeassistant.core import HomeAssistant
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.util import Throttle
from .const import (
@ -71,10 +70,8 @@ async def update_listener(hass: HomeAssistant, config_entry: ConfigEntry) -> Non
await hass.config_entries.async_reload(config_entry.entry_id)
registry = entity_registry.async_get(hass)
entities = entity_registry.async_entries_for_config_entry(
registry, config_entry.entry_id
)
registry = er.async_get(hass)
entities = er.async_entries_for_config_entry(registry, config_entry.entry_id)
# Remove orphaned entities
for entity in entities:

View file

@ -8,7 +8,11 @@ import coronavirus
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import aiohttp_client, entity_registry, update_coordinator
from homeassistant.helpers import (
aiohttp_client,
entity_registry as er,
update_coordinator,
)
from homeassistant.helpers.typing import ConfigType
from .const import DOMAIN
@ -31,16 +35,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
)
@callback
def _async_migrator(entity_entry: entity_registry.RegistryEntry):
def _async_migrator(entity_entry: er.RegistryEntry):
"""Migrate away from unstable ID."""
country, info_type = entity_entry.unique_id.rsplit("-", 1)
if not country.isnumeric():
return None
return {"new_unique_id": f"{entry.title}-{info_type}"}
await entity_registry.async_migrate_entries(
hass, entry.entry_id, _async_migrator
)
await er.async_migrate_entries(hass, entry.entry_id, _async_migrator)
if not entry.unique_id:
hass.config_entries.async_update_entry(entry, unique_id=entry.data["country"])

View file

@ -12,7 +12,7 @@ from homeassistant.components.device_tracker import (
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import STATE_UNKNOWN, UnitOfFrequency
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,
@ -30,7 +30,7 @@ async def async_setup_entry(
coordinators: dict[
str, DataUpdateCoordinator[list[ConnectedStationInfo]]
] = hass.data[DOMAIN][entry.entry_id]["coordinators"]
registry = entity_registry.async_get(hass)
registry = er.async_get(hass)
tracked = set()
@callback
@ -53,9 +53,7 @@ async def async_setup_entry(
def restore_entities() -> None:
"""Restore clients that are not a part of active clients list."""
missing = []
for entity in entity_registry.async_entries_for_config_entry(
registry, entry.entry_id
):
for entity in er.async_entries_for_config_entry(registry, entry.entry_id):
if (
entity.platform == DOMAIN
and entity.domain == DEVICE_TRACKER_DOMAIN

View file

@ -21,8 +21,7 @@ from homeassistant.const import CONF_DEVICE_ID, CONF_HOST, CONF_MAC, CONF_TYPE,
from homeassistant.core import HomeAssistant, callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.exceptions import IntegrationError
from homeassistant.helpers import device_registry
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import config_validation as cv, device_registry as dr
from .const import (
CONF_BROWSE_UNFILTERED,
@ -501,4 +500,4 @@ async def _async_get_mac_address(hass: HomeAssistant, host: str) -> str | None:
if not mac_address:
return None
return device_registry.format_mac(mac_address)
return dr.format_mac(mac_address)

View file

@ -29,7 +29,7 @@ from homeassistant.components.media_player import (
)
from homeassistant.const import CONF_DEVICE_ID, CONF_MAC, CONF_TYPE, CONF_URL
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.entity_platform import AddEntitiesCallback
from .const import (
@ -363,21 +363,21 @@ class DlnaDmrEntity(MediaPlayerEntity):
# device's UDN. They may be the same, if the DMR is the root device.
connections.add(
(
device_registry.CONNECTION_UPNP,
dr.CONNECTION_UPNP,
self._device.profile_device.root_device.udn,
)
)
connections.add((device_registry.CONNECTION_UPNP, self._device.udn))
connections.add((dr.CONNECTION_UPNP, self._device.udn))
if self.mac_address:
# Connection based on MAC address, if known
connections.add(
# Device MAC is obtained from the config entry, which uses getmac
(device_registry.CONNECTION_NETWORK_MAC, self.mac_address)
(dr.CONNECTION_NETWORK_MAC, self.mac_address)
)
# Create linked HA DeviceEntry now the information is known.
dev_reg = device_registry.async_get(self.hass)
dev_reg = dr.async_get(self.hass)
device_entry = dev_reg.async_get_or_create(
config_entry_id=self.registry_entry.config_entry_id,
connections=connections,
@ -388,7 +388,7 @@ class DlnaDmrEntity(MediaPlayerEntity):
)
# Update entity registry to link to the device
ent_reg = entity_registry.async_get(self.hass)
ent_reg = er.async_get(self.hass)
ent_reg.async_get_or_create(
self.registry_entry.domain,
self.registry_entry.platform,

View file

@ -9,8 +9,7 @@ import voluptuous as vol
from homeassistant.components.switch import PLATFORM_SCHEMA, SwitchEntity
from homeassistant.const import CONF_ID, CONF_NAME, Platform
from homeassistant.core import HomeAssistant
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.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
@ -38,7 +37,7 @@ def _migrate_to_new_unique_id(hass: HomeAssistant, dev_id, channel) -> None:
"""Migrate old unique ids to new unique ids."""
old_unique_id = f"{combine_hex(dev_id)}"
ent_reg = entity_registry.async_get(hass)
ent_reg = er.async_get(hass)
entity_id = ent_reg.async_get_entity_id(Platform.SWITCH, DOMAIN, old_unique_id)
if entity_id is not None:

View file

@ -3,7 +3,7 @@ from homeassistant.components.device_tracker import SourceType, TrackerEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_LATITUDE, 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.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -34,7 +34,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

@ -21,7 +21,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 import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
@ -257,7 +257,7 @@ async def async_setup_entry(
hass: HomeAssistant, old_unique_id: str, new_key: str
) -> None:
"""Migrate unique IDs to the new format."""
ent_reg = entity_registry.async_get(hass)
ent_reg = er.async_get(hass)
if entity_id := ent_reg.async_get_entity_id(
Platform.SENSOR, DOMAIN, old_unique_id

View file

@ -8,7 +8,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.dispatcher import async_dispatcher_connect
from homeassistant.helpers.entity import DeviceInfo
from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -44,7 +44,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

@ -12,7 +12,7 @@ from aioguardian.errors import GuardianError
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.update_coordinator import DataUpdateCoordinator, UpdateFailed
@ -41,7 +41,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

@ -5,7 +5,7 @@ from homeassistant.components.remote import ATTR_ACTIVITY, ATTR_DELAY_SECS
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_HOST, CONF_NAME, EVENT_HOMEASSISTANT_STOP
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_send
from .const import (
@ -60,7 +60,7 @@ async def _migrate_old_unique_ids(
names_to_ids = {activity["label"]: activity["id"] for activity in data.activities}
@callback
def _async_migrator(entity_entry: entity_registry.RegistryEntry):
def _async_migrator(entity_entry: er.RegistryEntry):
# Old format for switches was {remote_unique_id}-{activity_name}
# New format is activity_{activity_id}
parts = entity_entry.unique_id.split("-", 1)
@ -78,7 +78,7 @@ async def _migrate_old_unique_ids(
return None
await entity_registry.async_migrate_entries(hass, entry_id, _async_migrator)
await er.async_migrate_entries(hass, entry_id, _async_migrator)
@callback

View file

@ -50,8 +50,12 @@ from homeassistant.const import (
)
from homeassistant.core import CoreState, HomeAssistant, ServiceCall, State, callback
from homeassistant.exceptions import HomeAssistantError, Unauthorized
from homeassistant.helpers import device_registry, entity_registry, instance_id
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import (
config_validation as cv,
device_registry as dr,
entity_registry as er,
instance_id,
)
from homeassistant.helpers.entityfilter import (
BASE_FILTER_SCHEMA,
FILTER_SCHEMA,
@ -431,20 +435,19 @@ def _async_register_events_and_services(hass: HomeAssistant) -> None:
async def async_handle_homekit_unpair(service: ServiceCall) -> None:
"""Handle unpair HomeKit service call."""
referenced = async_extract_referenced_entity_ids(hass, service)
dev_reg = device_registry.async_get(hass)
dev_reg = dr.async_get(hass)
for device_id in referenced.referenced_devices:
if not (dev_reg_ent := dev_reg.async_get(device_id)):
raise HomeAssistantError(f"No device found for device id: {device_id}")
macs = [
cval
for ctype, cval in dev_reg_ent.connections
if ctype == device_registry.CONNECTION_NETWORK_MAC
if ctype == dr.CONNECTION_NETWORK_MAC
]
matching_instances = [
homekit
for homekit in _async_all_homekit_instances(hass)
if homekit.driver
and device_registry.format_mac(homekit.driver.state.mac) in macs
if homekit.driver and dr.format_mac(homekit.driver.state.mac) in macs
]
if not matching_instances:
raise HomeAssistantError(
@ -698,7 +701,7 @@ class HomeKit:
return False
def add_bridge_triggers_accessory(
self, device: device_registry.DeviceEntry, device_triggers: list[dict[str, Any]]
self, device: dr.DeviceEntry, device_triggers: list[dict[str, Any]]
) -> None:
"""Add device automation triggers to the bridge."""
if self._would_exceed_max_devices(device.name):
@ -734,8 +737,8 @@ class HomeKit:
async def async_configure_accessories(self) -> list[State]:
"""Configure accessories for the included states."""
dev_reg = device_registry.async_get(self.hass)
ent_reg = entity_registry.async_get(self.hass)
dev_reg = dr.async_get(self.hass)
ent_reg = er.async_get(self.hass)
device_lookup = ent_reg.async_get_device_class_lookup(
{
(BINARY_SENSOR_DOMAIN, BinarySensorDeviceClass.BATTERY_CHARGING),
@ -830,8 +833,8 @@ class HomeKit:
def _async_register_bridge(self) -> None:
"""Register the bridge as a device so homekit_controller and exclude it from discovery."""
assert self.driver is not None
dev_reg = device_registry.async_get(self.hass)
formatted_mac = device_registry.format_mac(self.driver.state.mac)
dev_reg = dr.async_get(self.hass)
formatted_mac = dr.format_mac(self.driver.state.mac)
# Connections and identifiers are both used here.
#
# connections exists so homekit_controller can know the
@ -844,7 +847,7 @@ class HomeKit:
# because this was the way you had to fix homekit when pairing
# failed.
#
connection = (device_registry.CONNECTION_NETWORK_MAC, formatted_mac)
connection = (dr.CONNECTION_NETWORK_MAC, formatted_mac)
identifier = (DOMAIN, self._entry_id, BRIDGE_SERIAL_NUMBER)
self._async_purge_old_bridges(dev_reg, identifier, connection)
is_accessory_mode = self._homekit_mode == HOMEKIT_MODE_ACCESSORY
@ -858,13 +861,13 @@ class HomeKit:
manufacturer=MANUFACTURER,
name=accessory_friendly_name(self._entry_title, self.driver.accessory),
model=f"HomeKit {hk_mode_name}",
entry_type=device_registry.DeviceEntryType.SERVICE,
entry_type=dr.DeviceEntryType.SERVICE,
)
@callback
def _async_purge_old_bridges(
self,
dev_reg: device_registry.DeviceRegistry,
dev_reg: dr.DeviceRegistry,
identifier: tuple[str, str, str],
connection: tuple[str, str],
) -> None:
@ -920,7 +923,7 @@ class HomeKit:
async def _async_add_trigger_accessories(self) -> None:
"""Add devices with triggers to the bridge."""
dev_reg = device_registry.async_get(self.hass)
dev_reg = dr.async_get(self.hass)
valid_device_ids = []
for device_id in self._devices:
if not dev_reg.async_get(device_id):
@ -989,7 +992,7 @@ class HomeKit:
@callback
def _async_configure_linked_sensors(
self,
ent_reg_ent: entity_registry.RegistryEntry,
ent_reg_ent: er.RegistryEntry,
device_lookup: dict[str, dict[tuple[str, str | None], str]],
state: State,
) -> None:
@ -1051,8 +1054,8 @@ class HomeKit:
async def _async_set_device_info_attributes(
self,
ent_reg_ent: entity_registry.RegistryEntry,
dev_reg: device_registry.DeviceRegistry,
ent_reg_ent: er.RegistryEntry,
dev_reg: dr.DeviceRegistry,
entity_id: str,
) -> None:
"""Set attributes that will be used for homekit device info."""
@ -1070,7 +1073,7 @@ class HomeKit:
ent_cfg[ATTR_INTEGRATION] = ent_reg_ent.platform
def _fill_config_from_device_registry_entry(
self, device_entry: device_registry.DeviceEntry, config: dict[str, Any]
self, device_entry: dr.DeviceEntry, config: dict[str, Any]
) -> None:
"""Populate a config dict from the registry."""
if device_entry.manufacturer:

View file

@ -28,8 +28,11 @@ from homeassistant.const import (
)
from homeassistant.core import HomeAssistant, callback, split_entity_id
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers import device_registry, entity_registry
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import (
config_validation as cv,
device_registry as dr,
entity_registry as er,
)
from homeassistant.helpers.entityfilter import (
CONF_EXCLUDE_DOMAINS,
CONF_EXCLUDE_ENTITIES,
@ -630,7 +633,7 @@ async def _async_get_supported_devices(hass: HomeAssistant) -> dict[str, str]:
results = await device_automation.async_get_device_automations(
hass, device_automation.DeviceAutomationType.TRIGGER
)
dev_reg = device_registry.async_get(hass)
dev_reg = dr.async_get(hass)
unsorted: dict[str, str] = {}
for device_id in results:
entry = dev_reg.async_get(device_id)
@ -639,7 +642,7 @@ async def _async_get_supported_devices(hass: HomeAssistant) -> dict[str, str]:
def _exclude_by_entity_registry(
ent_reg: entity_registry.EntityRegistry,
ent_reg: er.EntityRegistry,
entity_id: str,
include_entity_category: bool,
include_hidden: bool,
@ -661,7 +664,7 @@ def _async_get_matching_entities(
include_hidden: bool = False,
) -> dict[str, str]:
"""Fetch all entities or entities in the given domains."""
ent_reg = entity_registry.async_get(hass)
ent_reg = er.async_get(hass)
return {
state.entity_id: (
f"{state.attributes.get(ATTR_FRIENDLY_NAME, state.entity_id)} ({state.entity_id})"

View file

@ -7,7 +7,7 @@ from typing import Any
from pyhap.const import CATEGORY_SENSOR
from homeassistant.core import CALLBACK_TYPE, Context
from homeassistant.helpers import entity_registry
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.trigger import async_initialize_triggers
from .accessories import TYPES, HomeAccessory
@ -42,7 +42,7 @@ class DeviceTriggerAccessory(HomeAccessory):
self._remove_triggers: CALLBACK_TYPE | None = None
self.triggers = []
assert device_triggers is not None
ent_reg = entity_registry.async_get(self.hass)
ent_reg = er.async_get(self.hass)
for idx, trigger in enumerate(device_triggers):
type_: str = trigger["type"]
subtype: str | None = trigger.get("subtype")

View file

@ -44,7 +44,7 @@ from homeassistant.helpers import (
config_validation as cv,
device_registry as dr,
discovery,
entity_registry,
entity_registry as er,
)
from homeassistant.helpers.dispatcher import async_dispatcher_connect, dispatcher_send
from homeassistant.helpers.entity import DeviceInfo, Entity
@ -359,8 +359,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
# Transitional from < 2021.8: update None config entry and entity unique ids
if router_info and (serial_number := router_info.get("SerialNumber")):
hass.config_entries.async_update_entry(entry, unique_id=serial_number)
ent_reg = entity_registry.async_get(hass)
for entity_entry in entity_registry.async_entries_for_config_entry(
ent_reg = er.async_get(hass)
for entity_entry in er.async_entries_for_config_entry(
ent_reg, entry.entry_id
):
if not entity_entry.unique_id.startswith("None-"):

View file

@ -15,7 +15,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 import Entity
from homeassistant.helpers.entity_platform import AddEntitiesCallback
@ -66,7 +66,7 @@ async def async_setup_entry(
# Initialize already tracked entities
tracked: set[str] = set()
registry = entity_registry.async_get(hass)
registry = er.async_get(hass)
known_entities: list[Entity] = []
track_wired_clients = router.config_entry.options.get(
CONF_TRACK_WIRED_CLIENTS, DEFAULT_TRACK_WIRED_CLIENTS

View file

@ -18,8 +18,11 @@ from homeassistant.components import zeroconf
from homeassistant.const import CONF_API_KEY, CONF_HOST
from homeassistant.core import callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers import aiohttp_client, device_registry
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers import (
aiohttp_client,
config_validation as cv,
device_registry as dr,
)
from homeassistant.util.network import is_ipv6_address
from .const import (
@ -306,10 +309,8 @@ class HueV2OptionsFlowHandler(config_entries.OptionsFlow):
# create a list of Hue device ID's that the user can select
# to ignore availability status
dev_reg = device_registry.async_get(self.hass)
entries = device_registry.async_entries_for_config_entry(
dev_reg, self.config_entry.entry_id
)
dev_reg = dr.async_get(self.hass)
entries = dr.async_entries_for_config_entry(dev_reg, self.config_entry.entry_id)
dev_ids = {
identifier[1]: entry.name
for entry in entries

View file

@ -16,7 +16,7 @@ from homeassistant.const import (
ATTR_VIA_DEVICE,
)
from homeassistant.core import callback
from homeassistant.helpers import device_registry
from homeassistant.helpers import device_registry as dr
from ..const import DOMAIN
@ -29,11 +29,11 @@ async def async_setup_devices(bridge: "HueBridge"):
entry = bridge.config_entry
hass = bridge.hass
api: HueBridgeV2 = bridge.api # to satisfy typing
dev_reg = device_registry.async_get(hass)
dev_reg = dr.async_get(hass)
dev_controller = api.devices
@callback
def add_device(hue_device: Device) -> device_registry.DeviceEntry:
def add_device(hue_device: Device) -> dr.DeviceEntry:
"""Register a Hue device in device registry."""
model = f"{hue_device.product_data.product_name} ({hue_device.product_data.model_id})"
params = {
@ -51,9 +51,7 @@ async def async_setup_devices(bridge: "HueBridge"):
params[ATTR_VIA_DEVICE] = (DOMAIN, api.config.bridge_device.id)
zigbee = dev_controller.get_zigbee_connectivity(hue_device.id)
if zigbee and zigbee.mac_address:
params[ATTR_CONNECTIONS] = {
(device_registry.CONNECTION_NETWORK_MAC, zigbee.mac_address)
}
params[ATTR_CONNECTIONS] = {(dr.CONNECTION_NETWORK_MAC, zigbee.mac_address)}
return dev_reg.async_get_or_create(config_entry_id=entry.entry_id, **params)
@ -77,9 +75,7 @@ async def async_setup_devices(bridge: "HueBridge"):
known_devices = [add_device(hue_device) for hue_device in dev_controller]
# Check for nodes that no longer exist and remove them
for device in device_registry.async_entries_for_config_entry(
dev_reg, entry.entry_id
):
for device in dr.async_entries_for_config_entry(dev_reg, entry.entry_id):
if device not in known_devices:
# handle case where a virtual device was created for a Hue group
hue_dev_id = next(x[1] for x in device.identifiers if x[0] == DOMAIN)

View file

@ -9,7 +9,7 @@ from aiohue.v2.models.relative_rotary import RelativeRotary
from homeassistant.const import CONF_DEVICE_ID, CONF_ID, CONF_TYPE, CONF_UNIQUE_ID
from homeassistant.core import callback
from homeassistant.helpers import device_registry
from homeassistant.helpers import device_registry as dr
from homeassistant.util import slugify
from ..const import ATTR_HUE_EVENT, CONF_SUBTYPE, DOMAIN
@ -29,7 +29,7 @@ async def async_setup_hue_events(bridge: "HueBridge"):
hass = bridge.hass
api: HueBridgeV2 = bridge.api # to satisfy typing
conf_entry = bridge.config_entry
dev_reg = device_registry.async_get(hass)
dev_reg = dr.async_get(hass)
btn_controller = api.sensors.button
rotary_controller = api.sensors.relative_rotary