evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add support to entity registry for overriding device_class (#59985)

Browse source
This commit is contained in:
Erik Montnemery 2021-11-22 17:38:06 +01:00 • committed by GitHub
parent 67e13b35db
commit 49a27e12ad
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
31 changed files with 136 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

25
homeassistant/helpers/entity_registry.py
View file

@ -105,6 +105,7 @@ class RegistryEntry:
icon: str | None = attr.ib(default=None)
name: str | None = attr.ib(default=None)
# As set by integration
original_device_class: str | None = attr.ib(default=None)
original_icon: str | None = attr.ib(default=None)
original_name: str | None = attr.ib(default=None)
supported_features: int = attr.ib(default=0)
@ -128,8 +129,9 @@ class RegistryEntry:
if self.capabilities is not None:
attrs.update(self.capabilities)
if self.device_class is not None:
attrs[ATTR_DEVICE_CLASS] = self.device_class
device_class = self.device_class or self.original_device_class
if device_class is not None:
attrs[ATTR_DEVICE_CLASS] = device_class
icon = self.icon or self.original_icon
if icon is not None:
@ -190,7 +192,8 @@ class EntityRegistry:
for entity in self.entities.values():
if not entity.device_id:
continue
domain_device_class = (entity.domain, entity.device_class)
device_class = entity.device_class or entity.original_device_class
domain_device_class = (entity.domain, device_class)
if domain_device_class not in domain_device_classes:
continue
if entity.device_id not in lookup:
@ -268,9 +271,9 @@ class EntityRegistry:
area_id: str | None = None,
capabilities: Mapping[str, Any] | None = None,
config_entry: ConfigEntry | None = None,
device_class: str | None = None,
device_id: str | None = None,
entity_category: str | None = None,
original_device_class: str | None = None,
original_icon: str | None = None,
original_name: str | None = None,
supported_features: int | None = None,
@ -289,9 +292,9 @@ class EntityRegistry:
area_id=area_id or UNDEFINED,
capabilities=capabilities or UNDEFINED,
config_entry_id=config_entry_id or UNDEFINED,
device_class=device_class or UNDEFINED,
device_id=device_id or UNDEFINED,
entity_category=entity_category or UNDEFINED,
original_device_class=original_device_class or UNDEFINED,
original_icon=original_icon or UNDEFINED,
original_name=original_name or UNDEFINED,
supported_features=supported_features or UNDEFINED,
@ -320,11 +323,11 @@ class EntityRegistry:
area_id=area_id,
capabilities=capabilities,
config_entry_id=config_entry_id,
device_class=device_class,
device_id=device_id,
disabled_by=disabled_by,
entity_category=entity_category,
entity_id=entity_id,
original_device_class=original_device_class,
original_icon=original_icon,
original_name=original_name,
platform=platform,
@ -409,6 +412,7 @@ class EntityRegistry:
name: str | None | UndefinedType = UNDEFINED,
new_entity_id: str | UndefinedType = UNDEFINED,
new_unique_id: str | UndefinedType = UNDEFINED,
original_device_class: str | None | UndefinedType = UNDEFINED,
original_icon: str | None | UndefinedType = UNDEFINED,
original_name: str | None | UndefinedType = UNDEFINED,
unit_of_measurement: str | None | UndefinedType = UNDEFINED,
@ -425,6 +429,7 @@ class EntityRegistry:
name=name,
new_entity_id=new_entity_id,
new_unique_id=new_unique_id,
original_device_class=original_device_class,
original_icon=original_icon,
original_name=original_name,
unit_of_measurement=unit_of_measurement,
@ -446,6 +451,7 @@ class EntityRegistry:
name: str | None | UndefinedType = UNDEFINED,
new_entity_id: str | UndefinedType = UNDEFINED,
new_unique_id: str | UndefinedType = UNDEFINED,
original_device_class: str | None | UndefinedType = UNDEFINED,
original_icon: str | None | UndefinedType = UNDEFINED,
original_name: str | None | UndefinedType = UNDEFINED,
supported_features: int | UndefinedType = UNDEFINED,
@ -467,6 +473,7 @@ class EntityRegistry:
("entity_category", entity_category),
("icon", icon),
("name", name),
("original_device_class", original_device_class),
("original_icon", original_icon),
("original_name", original_name),
("supported_features", supported_features),
@ -552,6 +559,7 @@ class EntityRegistry:
entity_id=entity["entity_id"],
icon=entity["icon"],
name=entity["name"],
original_device_class=entity["original_device_class"],
original_icon=entity["original_icon"],
original_name=entity["original_name"],
platform=entity["platform"],
@ -585,6 +593,7 @@ class EntityRegistry:
"entity_id": entry.entity_id,
"icon": entry.icon,
"name": entry.name,
"original_device_class": entry.original_device_class,
"original_icon": entry.original_icon,
"original_name": entry.original_name,
"platform": entry.platform,
@ -741,6 +750,10 @@ async def _async_migrate(
entity["platform"] = entity["platform"]
entity["supported_features"] = entity.get("supported_features", 0)
entity["unit_of_measurement"] = entity.get("unit_of_measurement")
# Move device_class to original_device_class
entity["original_device_class"] = entity["device_class"]
entity["device_class"] = None
if old_major_version > 1:
raise NotImplementedError
return data