Jinja filter and function for median and statistical_mode (#105554)

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>
Co-authored-by: Franck Nijhof <git@frenck.dev>

homeassistant/helpers/template.py

@@ -1908,6 +1908,66 @@ def average(*args: Any, default: Any = _SENTINEL) -> Any:
         return default
 
 
+def median(*args: Any, default: Any = _SENTINEL) -> Any:
+    """Filter and function to calculate the median.
+
+    Calculates median of an iterable of two or more arguments.
+
+    The parameters may be passed as an iterable or as separate arguments.
+    """
+    if len(args) == 0:
+        raise TypeError("median expected at least 1 argument, got 0")
+
+    # If first argument is a list or tuple and more than 1 argument provided but not a named
+    # default, then use 2nd argument as default.
+    if isinstance(args[0], Iterable):
+        median_list = args[0]
+        if len(args) > 1 and default is _SENTINEL:
+            default = args[1]
+    elif len(args) == 1:
+        raise TypeError(f"'{type(args[0]).__name__}' object is not iterable")
+    else:
+        median_list = args
+
+    try:
+        return statistics.median(median_list)
+    except (TypeError, statistics.StatisticsError):
+        if default is _SENTINEL:
+            raise_no_default("median", args)
+        return default
+
+
+def statistical_mode(*args: Any, default: Any = _SENTINEL) -> Any:
+    """Filter and function to calculate the statistical mode.
+
+    Calculates mode of an iterable of two or more arguments.
+
+    The parameters may be passed as an iterable or as separate arguments.
+    """
+    if not args:
+        raise TypeError("statistical_mode expected at least 1 argument, got 0")
+
+    # If first argument is a list or tuple and more than 1 argument provided but not a named
+    # default, then use 2nd argument as default.
+    if len(args) == 1 and isinstance(args[0], Iterable):
+        mode_list = args[0]
+    elif isinstance(args[0], list | tuple):
+        mode_list = args[0]
+        if len(args) > 1 and default is _SENTINEL:
+            default = args[1]
+    elif len(args) == 1:
+        raise TypeError(f"'{type(args[0]).__name__}' object is not iterable")
+    else:
+        mode_list = args
+
+    try:
+        return statistics.mode(mode_list)
+    except (TypeError, statistics.StatisticsError):
+        if default is _SENTINEL:
+            raise_no_default("statistical_mode", args)
+        return default
+
+
 def forgiving_float(value, default=_SENTINEL):
     """Try to convert value to a float."""
     try:
@@ -2390,6 +2450,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         self.filters["from_json"] = from_json
         self.filters["is_defined"] = fail_when_undefined
         self.filters["average"] = average
+        self.filters["median"] = median
+        self.filters["statistical_mode"] = statistical_mode
         self.filters["random"] = random_every_time
         self.filters["base64_encode"] = base64_encode
         self.filters["base64_decode"] = base64_decode
@@ -2412,6 +2474,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         self.filters["bool"] = forgiving_boolean
         self.filters["version"] = version
         self.filters["contains"] = contains
+        self.filters["median"] = median
+        self.filters["statistical_mode"] = statistical_mode
         self.globals["log"] = logarithm
         self.globals["sin"] = sine
         self.globals["cos"] = cosine
@@ -2433,6 +2497,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         self.globals["strptime"] = strptime
         self.globals["urlencode"] = urlencode
         self.globals["average"] = average
+        self.globals["median"] = median
+        self.globals["statistical_mode"] = statistical_mode
         self.globals["max"] = min_max_from_filter(self.filters["max"], "max")
         self.globals["min"] = min_max_from_filter(self.filters["min"], "min")
         self.globals["is_number"] = is_number
@@ -2445,6 +2511,8 @@ class TemplateEnvironment(ImmutableSandboxedEnvironment):
         self.globals["iif"] = iif
         self.globals["bool"] = forgiving_boolean
         self.globals["version"] = version
+        self.globals["median"] = median
+        self.globals["statistical_mode"] = statistical_mode
         self.tests["is_number"] = is_number
         self.tests["list"] = _is_list
         self.tests["set"] = _is_set

tests/helpers/test_template.py

@@ -1318,6 +1318,88 @@ def test_average(hass: HomeAssistant) -> None:
         template.Template("{{ average([]) }}", hass).async_render()
 
 
+def test_median(hass: HomeAssistant) -> None:
+    """Test the median filter."""
+    assert template.Template("{{ [1, 3, 2] | median }}", hass).async_render() == 2
+    assert template.Template("{{ median([1, 3, 2, 4]) }}", hass).async_render() == 2.5
+    assert template.Template("{{ median(1, 3, 2) }}", hass).async_render() == 2
+    assert template.Template("{{ median('cdeba') }}", hass).async_render() == "c"
+
+    # Testing of default values
+    assert template.Template("{{ median([1, 2, 3], -1) }}", hass).async_render() == 2
+    assert template.Template("{{ median([], -1) }}", hass).async_render() == -1
+    assert template.Template("{{ median([], default=-1) }}", hass).async_render() == -1
+    assert template.Template("{{ median('abcd', -1) }}", hass).async_render() == -1
+    assert (
+        template.Template("{{ median([], 5, default=-1) }}", hass).async_render() == -1
+    )
+    assert (
+        template.Template("{{ median(1, 'a', 3, default=-1) }}", hass).async_render()
+        == -1
+    )
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ 1 | median }}", hass).async_render()
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ median() }}", hass).async_render()
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ median([]) }}", hass).async_render()
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ median('abcd') }}", hass).async_render()
+
+
+def test_statistical_mode(hass: HomeAssistant) -> None:
+    """Test the mode filter."""
+    assert (
+        template.Template("{{ [1, 2, 2, 3] | statistical_mode }}", hass).async_render()
+        == 2
+    )
+    assert (
+        template.Template("{{ statistical_mode([1, 2, 3]) }}", hass).async_render() == 1
+    )
+    assert (
+        template.Template(
+            "{{ statistical_mode('hello', 'bye', 'hello') }}", hass
+        ).async_render()
+        == "hello"
+    )
+    assert (
+        template.Template("{{ statistical_mode('banana') }}", hass).async_render()
+        == "a"
+    )
+
+    # Testing of default values
+    assert (
+        template.Template("{{ statistical_mode([1, 2, 3], -1) }}", hass).async_render()
+        == 1
+    )
+    assert (
+        template.Template("{{ statistical_mode([], -1) }}", hass).async_render() == -1
+    )
+    assert (
+        template.Template("{{ statistical_mode([], default=-1) }}", hass).async_render()
+        == -1
+    )
+    assert (
+        template.Template(
+            "{{ statistical_mode([], 5, default=-1) }}", hass
+        ).async_render()
+        == -1
+    )
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ 1 | statistical_mode }}", hass).async_render()
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ statistical_mode() }}", hass).async_render()
+
+    with pytest.raises(TemplateError):
+        template.Template("{{ statistical_mode([]) }}", hass).async_render()
+
+
 def test_min(hass: HomeAssistant) -> None:
     """Test the min filter."""
     assert template.Template("{{ [1, 2, 3] | min }}", hass).async_render() == 1