refactor: simplify entity attribute retrieval using state_attr (#162)

hugobloem · Dec 24, 2024 · aa0eba1 · aa0eba1
1 parent 7c984b2
commit aa0eba1
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/custom_components/stateful_scenes/helpers.py b/custom_components/stateful_scenes/helpers.py
@@ -2,6 +2,7 @@
 
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry, device_registry, area_registry
+from homeassistant.helpers.template import state_attr
 
 
 def get_id_from_entity_id(hass: HomeAssistant, entity_id: str) -> str:
@@ -12,16 +13,12 @@ def get_id_from_entity_id(hass: HomeAssistant, entity_id: str) -> str:
 
 def get_name_from_entity_id(hass: HomeAssistant, entity_id: str) -> str:
     """Get scene name from entity_id."""
-    er = entity_registry.async_get(hass)
-    name = er.async_get(entity_id).original_name
-    return name if name is not None else entity_id
+    return state_attr(hass, entity_id, "friendly_name")
 
 
 def get_icon_from_entity_id(hass: HomeAssistant, entity_id: str) -> str:
     """Get scene icon from entity_id."""
-    er = entity_registry.async_get(hass)
-    icon = er.async_get(entity_id).icon
-    return icon if icon is not None else None
+    return state_attr(hass, entity_id, "icon")
 
 
 def get_area_from_entity_id(hass: HomeAssistant, entity_id: str) -> str: