Add template function: flatten (#140157)

This commit is contained in:
Franck Nijhof 2025-03-08 22:14:27 +01:00 committed by GitHub
parent 62c025fd12
commit 9aa8a786a5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 58 additions and 0 deletions

View File

@ -2765,6 +2765,25 @@ def typeof(value: Any) -> Any:
return value.__class__.__name__
def flatten(value: Iterable[Any], levels: int | None = None) -> list[Any]:
"""Flattens list of lists."""
if not isinstance(value, Iterable) or isinstance(value, str):
raise TypeError(f"flatten expected a list, got {type(value).__name__}")
flattened: list[Any] = []
for item in value:
if isinstance(item, Iterable) and not isinstance(item, str):
if levels is None:
flattened.extend(flatten(item))
elif levels >= 1:
flattened.extend(flatten(item, levels=(levels - 1)))
else:
flattened.append(item)
else:
flattened.append(item)
return flattened
class TemplateContextManager(AbstractContextManager):
"""Context manager to store template being parsed or rendered in a ContextVar."""
@ -2967,6 +2986,7 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
self.filters["contains"] = contains
self.filters["shuffle"] = shuffle
self.filters["typeof"] = typeof
self.filters["flatten"] = flatten
self.globals["log"] = logarithm
self.globals["sin"] = sine
self.globals["cos"] = cosine
@ -3006,6 +3026,7 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
self.globals["zip"] = zip
self.globals["shuffle"] = shuffle
self.globals["typeof"] = typeof
self.globals["flatten"] = flatten
self.tests["is_number"] = is_number
self.tests["list"] = _is_list
self.tests["set"] = _is_set

View File

@ -6751,3 +6751,40 @@ def test_typeof(hass: HomeAssistant) -> None:
template.Template("{{ typeof('Home Assistant') }}", hass).async_render()
== "str"
)
def test_flatten(hass: HomeAssistant) -> None:
"""Test the flatten function and filter."""
assert template.Template(
"{{ flatten([1, [2, [3]], 4, [5 , 6]]) }}", hass
).async_render() == [1, 2, 3, 4, 5, 6]
assert template.Template(
"{{ [1, [2, [3]], 4, [5 , 6]] | flatten }}", hass
).async_render() == [1, 2, 3, 4, 5, 6]
assert template.Template(
"{{ flatten([1, [2, [3]], 4, [5 , 6]], 1) }}", hass
).async_render() == [1, 2, [3], 4, 5, 6]
assert template.Template(
"{{ flatten([1, [2, [3]], 4, [5 , 6]], levels=1) }}", hass
).async_render() == [1, 2, [3], 4, 5, 6]
assert template.Template(
"{{ [1, [2, [3]], 4, [5 , 6]] | flatten(1) }}", hass
).async_render() == [1, 2, [3], 4, 5, 6]
assert template.Template(
"{{ [1, [2, [3]], 4, [5 , 6]] | flatten(levels=1) }}", hass
).async_render() == [1, 2, [3], 4, 5, 6]
assert template.Template("{{ flatten([]) }}", hass).async_render() == []
assert template.Template("{{ [] | flatten }}", hass).async_render() == []
with pytest.raises(TemplateError):
template.Template("{{ 'string' | flatten }}", hass).async_render()
with pytest.raises(TemplateError):
template.Template("{{ flatten() }}", hass).async_render()