Add floor template functions (#110847)

homeassistant/helpers/template.py

@@ -83,6 +83,7 @@ from . import (
     area_registry,
     device_registry,
     entity_registry,
+    floor_registry as fr,
     issue_registry,
     location as loc_helper,
 )
@@ -1387,6 +1388,45 @@ 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 name."""
+    floor_registry = fr.async_get(hass)
+    if floor := floor_registry.async_get_floor_by_name(str(lookup_value)):
+        return floor.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
+    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 = area_registry.async_get(hass)
+    entries = area_registry.async_entries_for_floor(area_reg, _floor_id)
+    return [entry.id for entry in entries if entry.id]
+
+
 def areas(hass: HomeAssistant) -> Iterable[str | None]:
     """Return all areas."""
     area_reg = area_registry.async_get(hass)
@@ -2668,6 +2708,18 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         self.globals["area_devices"] = hassfunction(area_devices)
         self.filters["area_devices"] = self.globals["area_devices"]
 
+        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["integration_entities"] = hassfunction(integration_entities)
         self.filters["integration_entities"] = self.globals["integration_entities"]
 
@@ -2700,6 +2752,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
                 "device_id",
                 "area_id",
                 "area_name",
+                "floor_id",
+                "floor_name",
                 "relative_time",
                 "today_at",
             ]
@@ -2709,6 +2763,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
                 "device_id",
                 "area_id",
                 "area_name",
+                "floor_id",
+                "floor_name",
                 "has_value",
             ]
             hass_tests = [

tests/helpers/test_template.py

@@ -38,6 +38,7 @@ from homeassistant.helpers import (
     device_registry as dr,
     entity,
     entity_registry as er,
+    floor_registry as fr,
     issue_registry as ir,
     template,
     translation,
@@ -5167,3 +5168,144 @@ async def test_lru_increases_with_many_entities(hass: HomeAssistant) -> None:
     assert template.CACHED_TEMPLATE_NO_COLLECT_LRU.get_size() == int(
         round(mock_entity_count * template.ENTITY_COUNT_GROWTH_FACTOR)
     )
+
+
+async def test_floors(
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
+) -> None:
+    """Test floors function."""
+
+    # Test no floors
+    info = render_to_info(hass, "{{ floors() }}")
+    assert_result_info(info, [])
+    assert info.rate_limit is None
+
+    # Test one floor
+    floor1 = floor_registry.async_create("First floor")
+    info = render_to_info(hass, "{{ floors() }}")
+    assert_result_info(info, [floor1.floor_id])
+    assert info.rate_limit is None
+
+    # Test multiple floors
+    floor2 = floor_registry.async_create("Second floor")
+    info = render_to_info(hass, "{{ floors() }}")
+    assert_result_info(info, [floor1.floor_id, floor2.floor_id])
+    assert info.rate_limit is None
+
+
+async def test_floor_id(
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
+) -> None:
+    """Test floor_id function."""
+
+    # Test non existing floor name
+    info = render_to_info(hass, "{{ floor_id('Third floor') }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 'Third floor' | floor_id }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    # Test wrong value type
+    info = render_to_info(hass, "{{ floor_id(42) }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 42 | floor_id }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    # Test with an actual floor
+    floor = floor_registry.async_create("First floor")
+    info = render_to_info(hass, "{{ floor_id('First floor') }}")
+    assert_result_info(info, floor.floor_id)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 'First floor' | floor_id }}")
+    assert_result_info(info, floor.floor_id)
+    assert info.rate_limit is None
+
+
+async def test_floor_name(
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
+) -> None:
+    """Test floor_name function."""
+    # Test non existing floor ID
+    info = render_to_info(hass, "{{ floor_name('third_floor') }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 'third_floor' | floor_name }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    # Test wrong value type
+    info = render_to_info(hass, "{{ floor_name(42) }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 42 | floor_name }}")
+    assert_result_info(info, None)
+    assert info.rate_limit is None
+
+    # Test existing floor ID
+    floor = floor_registry.async_create("First floor")
+    info = render_to_info(hass, f"{{{{ floor_name('{floor.floor_id}') }}}}")
+    assert_result_info(info, floor.name)
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, f"{{{{ '{floor.floor_id}' | floor_name }}}}")
+    assert_result_info(info, floor.name)
+    assert info.rate_limit is None
+
+
+async def test_floor_areas(
+    hass: HomeAssistant,
+    floor_registry: fr.FloorRegistry,
+    area_registry: ar.AreaRegistry,
+) -> None:
+    """Test floor_areas function."""
+
+    # Test non existing floor ID
+    info = render_to_info(hass, "{{ floor_areas('skyring') }}")
+    assert_result_info(info, [])
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 'skyring' | floor_areas }}")
+    assert_result_info(info, [])
+    assert info.rate_limit is None
+
+    # Test wrong value type
+    info = render_to_info(hass, "{{ floor_areas(42) }}")
+    assert_result_info(info, [])
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, "{{ 42 | floor_areas }}")
+    assert_result_info(info, [])
+    assert info.rate_limit is None
+
+    floor = floor_registry.async_create("First floor")
+    area = area_registry.async_create("Living room")
+    area_registry.async_update(area.id, floor_id=floor.floor_id)
+
+    # Get areas by floor ID
+    info = render_to_info(hass, f"{{{{ floor_areas('{floor.floor_id}') }}}}")
+    assert_result_info(info, [area.id])
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, f"{{{{ '{floor.floor_id}' | floor_areas }}}}")
+    assert_result_info(info, [area.id])
+    assert info.rate_limit is None
+
+    # Get entities by floor name
+    info = render_to_info(hass, f"{{{{ floor_areas('{floor.name}') }}}}")
+    assert_result_info(info, [area.id])
+    assert info.rate_limit is None
+
+    info = render_to_info(hass, f"{{{{ '{floor.name}' | floor_areas }}}}")
+    assert_result_info(info, [area.id])
+    assert info.rate_limit is None