Extract floor template functions into a floors Jinja2 extension (home-assistant#156589)

Co-authored-by: Copilot <[email protected]>

Author: frenck

homeassistant/helpers/template/__init__.py

@@ -59,7 +59,6 @@
     area_registry as ar,
     device_registry as dr,
     entity_registry as er,
-    floor_registry as fr,
     issue_registry as ir,
     location as loc_helper,
 )
@@ -79,7 +78,7 @@
     template_context_manager,
     template_cv,
 )
-from .helpers import raise_no_default
+from .helpers import raise_no_default, resolve_area_id
 from .render_info import RenderInfo, render_info_cv
 
 if TYPE_CHECKING:
@@ -1318,112 +1317,14 @@ def issue(hass: HomeAssistant, domain: str, issue_id: str) -> dict[str, Any] | N
     return None
 
 
-def floors(hass: HomeAssistant) -> Iterable[str | None]:
-    """Return all floors."""
-    floor_registry = fr.async_get(hass)
-    return [floor.floor_id for floor in floor_registry.async_list_floors()]
-
-
-def floor_id(hass: HomeAssistant, lookup_value: Any) -> str | None:
-    """Get the floor ID from a floor or area name, alias, device id, or entity id."""
-    floor_registry = fr.async_get(hass)
-    lookup_str = str(lookup_value)
-    if floor := floor_registry.async_get_floor_by_name(lookup_str):
-        return floor.floor_id
-    floors_list = floor_registry.async_get_floors_by_alias(lookup_str)
-    if floors_list:
-        return floors_list[0].floor_id
-
-    if aid := area_id(hass, lookup_value):
-        area_reg = ar.async_get(hass)
-        if area := area_reg.async_get_area(aid):
-            return area.floor_id
-
-    return None
-
-
-def floor_name(hass: HomeAssistant, lookup_value: str) -> str | None:
-    """Get the floor name from a floor id."""
-    floor_registry = fr.async_get(hass)
-    if floor := floor_registry.async_get_floor(lookup_value):
-        return floor.name
-
-    if aid := area_id(hass, lookup_value):
-        area_reg = ar.async_get(hass)
-        if (
-            (area := area_reg.async_get_area(aid))
-            and area.floor_id
-            and (floor := floor_registry.async_get_floor(area.floor_id))
-        ):
-            return floor.name
-
-    return None
-
-
-def floor_areas(hass: HomeAssistant, floor_id_or_name: str) -> Iterable[str]:
-    """Return area IDs for a given floor ID or name."""
-    _floor_id: str | None
-    # If floor_name returns a value, we know the input was an ID, otherwise we
-    # assume it's a name, and if it's neither, we return early
-    if floor_name(hass, floor_id_or_name) is not None:
-        _floor_id = floor_id_or_name
-    else:
-        _floor_id = floor_id(hass, floor_id_or_name)
-    if _floor_id is None:
-        return []
-
-    area_reg = ar.async_get(hass)
-    entries = ar.async_entries_for_floor(area_reg, _floor_id)
-    return [entry.id for entry in entries if entry.id]
-
-
-def floor_entities(hass: HomeAssistant, floor_id_or_name: str) -> Iterable[str]:
-    """Return entity_ids for a given floor ID or name."""
-    return [
-        entity_id
-        for area_id in floor_areas(hass, floor_id_or_name)
-        for entity_id in area_entities(hass, area_id)
-    ]
-
-
 def areas(hass: HomeAssistant) -> Iterable[str | None]:
     """Return all areas."""
     return list(ar.async_get(hass).areas)
 
 
 def area_id(hass: HomeAssistant, lookup_value: str) -> str | None:
     """Get the area ID from an area name, alias, device id, or entity id."""
-    area_reg = ar.async_get(hass)
-    lookup_str = str(lookup_value)
-    if area := area_reg.async_get_area_by_name(lookup_str):
-        return area.id
-    areas_list = area_reg.async_get_areas_by_alias(lookup_str)
-    if areas_list:
-        return areas_list[0].id
-
-    ent_reg = er.async_get(hass)
-    dev_reg = dr.async_get(hass)
-    # Import here, not at top-level to avoid circular import
-    from homeassistant.helpers import config_validation as cv  # noqa: PLC0415
-
-    try:
-        cv.entity_id(lookup_value)
-    except vol.Invalid:
-        pass
-    else:
-        if entity := ent_reg.async_get(lookup_value):
-            # If entity has an area ID, return that
-            if entity.area_id:
-                return entity.area_id
-            # If entity has a device ID, return the area ID for the device
-            if entity.device_id and (device := dev_reg.async_get(entity.device_id)):
-                return device.area_id
-
-    # Check if this could be a device ID
-    if device := dev_reg.async_get(lookup_value):
-        return device.area_id
-
-    return None
+    return resolve_area_id(hass, lookup_value)
 
 
 def _get_area_name(area_reg: ar.AreaRegistry, valid_area_id: str) -> str:
@@ -2359,6 +2260,7 @@ def __init__(
             "homeassistant.helpers.template.extensions.CollectionExtension"
         )
         self.add_extension("homeassistant.helpers.template.extensions.CryptoExtension")
+        self.add_extension("homeassistant.helpers.template.extensions.FloorExtension")
         self.add_extension("homeassistant.helpers.template.extensions.LabelExtension")
         self.add_extension("homeassistant.helpers.template.extensions.MathExtension")
         self.add_extension("homeassistant.helpers.template.extensions.RegexExtension")
@@ -2462,23 +2364,6 @@ def wrapper(_: Any, *args: _P.args, **kwargs: _P.kwargs) -> _R:
         self.globals["area_devices"] = hassfunction(area_devices)
         self.filters["area_devices"] = self.globals["area_devices"]
 
-        # Floor extensions
-
-        self.globals["floors"] = hassfunction(floors)
-        self.filters["floors"] = self.globals["floors"]
-
-        self.globals["floor_id"] = hassfunction(floor_id)
-        self.filters["floor_id"] = self.globals["floor_id"]
-
-        self.globals["floor_name"] = hassfunction(floor_name)
-        self.filters["floor_name"] = self.globals["floor_name"]
-
-        self.globals["floor_areas"] = hassfunction(floor_areas)
-        self.filters["floor_areas"] = self.globals["floor_areas"]
-
-        self.globals["floor_entities"] = hassfunction(floor_entities)
-        self.filters["floor_entities"] = self.globals["floor_entities"]
-
         # Integration extensions
 
         self.globals["integration_entities"] = hassfunction(integration_entities)
@@ -2534,8 +2419,6 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "device_id",
                 "distance",
                 "expand",
-                "floor_id",
-                "floor_name",
                 "has_value",
                 "is_device_attr",
                 "is_hidden_entity",
@@ -2557,8 +2440,6 @@ def warn_unsupported(*args: Any, **kwargs: Any) -> NoReturn:
                 "closest",
                 "device_id",
                 "expand",
-                "floor_id",
-                "floor_name",
                 "has_value",
             ]
             hass_tests = [

homeassistant/helpers/template/extensions/__init__.py

@@ -3,6 +3,7 @@
 from .base64 import Base64Extension
 from .collection import CollectionExtension
 from .crypto import CryptoExtension
+from .floors import FloorExtension
 from .labels import LabelExtension
 from .math import MathExtension
 from .regex import RegexExtension
@@ -12,6 +13,7 @@
     "Base64Extension",
     "CollectionExtension",
     "CryptoExtension",
+    "FloorExtension",
     "LabelExtension",
     "MathExtension",
     "RegexExtension",

homeassistant/helpers/template/extensions/floors.py

@@ -0,0 +1,157 @@
+"""Floor functions for Home Assistant templates."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any
+
+from homeassistant.helpers import (
+    area_registry as ar,
+    device_registry as dr,
+    entity_registry as er,
+    floor_registry as fr,
+)
+from homeassistant.helpers.template.helpers import resolve_area_id
+
+from .base import BaseTemplateExtension, TemplateFunction
+
+if TYPE_CHECKING:
+    from homeassistant.helpers.template import TemplateEnvironment
+
+
+class FloorExtension(BaseTemplateExtension):
+    """Extension for floor-related template functions."""
+
+    def __init__(self, environment: TemplateEnvironment) -> None:
+        """Initialize the floor extension."""
+        super().__init__(
+            environment,
+            functions=[
+                TemplateFunction(
+                    "floors",
+                    self.floors,
+                    as_global=True,
+                    requires_hass=True,
+                ),
+                TemplateFunction(
+                    "floor_id",
+                    self.floor_id,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+                TemplateFunction(
+                    "floor_name",
+                    self.floor_name,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                    limited_ok=False,
+                ),
+                TemplateFunction(
+                    "floor_areas",
+                    self.floor_areas,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                ),
+                TemplateFunction(
+                    "floor_entities",
+                    self.floor_entities,
+                    as_global=True,
+                    as_filter=True,
+                    requires_hass=True,
+                ),
+            ],
+        )
+
+    def floors(self) -> Iterable[str | None]:
+        """Return all floors."""
+        floor_registry = fr.async_get(self.hass)
+        return [floor.floor_id for floor in floor_registry.async_list_floors()]
+
+    def floor_id(self, lookup_value: Any) -> str | None:
+        """Get the floor ID from a floor or area name, alias, device id, or entity id."""
+        floor_registry = fr.async_get(self.hass)
+        lookup_str = str(lookup_value)
+
+        # Check if it's a floor name or alias
+        if floor := floor_registry.async_get_floor_by_name(lookup_str):
+            return floor.floor_id
+        floors_list = floor_registry.async_get_floors_by_alias(lookup_str)
+        if floors_list:
+            return floors_list[0].floor_id
+
+        # Resolve to area ID and get floor from area
+        if aid := resolve_area_id(self.hass, lookup_value):
+            area_reg = ar.async_get(self.hass)
+            if area := area_reg.async_get_area(aid):
+                return area.floor_id
+
+        return None
+
+    def floor_name(self, lookup_value: str) -> str | None:
+        """Get the floor name from a floor id."""
+        floor_registry = fr.async_get(self.hass)
+
+        # Check if it's a floor ID
+        if floor := floor_registry.async_get_floor(lookup_value):
+            return floor.name
+
+        # Resolve to area ID and get floor name from area's floor
+        if aid := resolve_area_id(self.hass, lookup_value):
+            area_reg = ar.async_get(self.hass)
+            if (
+                (area := area_reg.async_get_area(aid))
+                and area.floor_id
+                and (floor := floor_registry.async_get_floor(area.floor_id))
+            ):
+                return floor.name
+
+        return None
+
+    def _floor_id_or_name(self, floor_id_or_name: str) -> str | None:
+        """Get the floor ID from a floor name or ID."""
+        # If floor_name returns a value, we know the input was an ID, otherwise we
+        # assume it's a name, and if it's neither, we return early.
+        if self.floor_name(floor_id_or_name) is not None:
+            return floor_id_or_name
+        return self.floor_id(floor_id_or_name)
+
+    def floor_areas(self, floor_id_or_name: str) -> Iterable[str]:
+        """Return area IDs for a given floor ID or name."""
+        if (_floor_id := self._floor_id_or_name(floor_id_or_name)) is None:
+            return []
+
+        area_reg = ar.async_get(self.hass)
+        entries = ar.async_entries_for_floor(area_reg, _floor_id)
+        return [entry.id for entry in entries if entry.id]
+
+    def floor_entities(self, floor_id_or_name: str) -> Iterable[str]:
+        """Return entity_ids for a given floor ID or name."""
+        ent_reg = er.async_get(self.hass)
+        dev_reg = dr.async_get(self.hass)
+        entity_ids = []
+
+        for area_id in self.floor_areas(floor_id_or_name):
+            # Get entities directly assigned to the area
+            entity_ids.extend(
+                [
+                    entry.entity_id
+                    for entry in er.async_entries_for_area(ent_reg, area_id)
+                ]
+            )
+
+            # Also add entities tied to a device in the area that don't themselves
+            # have an area specified since they inherit the area from the device
+            entity_ids.extend(
+                [
+                    entity.entity_id
+                    for device in dr.async_entries_for_area(dev_reg, area_id)
+                    for entity in er.async_entries_for_device(ent_reg, device.id)
+                    if entity.area_id is None
+                ]
+            )
+
+        return entity_ids