Refactor template engine: Extract math & statistical functions into MathExtension (home-assistant#152266)

Author: frenck

homeassistant/helpers/template/__init__.py

@@ -18,7 +18,6 @@
 import pathlib
 import random
 import re
-import statistics
 from struct import error as StructError, pack, unpack_from
 import sys
 from types import CodeType, TracebackType
@@ -37,7 +36,7 @@
 
 from awesomeversion import AwesomeVersion
 import jinja2
-from jinja2 import pass_context, pass_environment, pass_eval_context
+from jinja2 import pass_context, pass_eval_context
 from jinja2.runtime import AsyncLoopContext, LoopContext
 from jinja2.sandbox import ImmutableSandboxedEnvironment
 from jinja2.utils import Namespace
@@ -2047,121 +2046,11 @@ def returns(value):
     return wrapper
 
 
-def logarithm(value, base=math.e, default=_SENTINEL):
-    """Filter and function to get logarithm of the value with a specific base."""
-    try:
-        base_float = float(base)
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("log", base)
-        return default
-    try:
-        value_float = float(value)
-        return math.log(value_float, base_float)
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("log", value)
-        return default
-
-
-def sine(value, default=_SENTINEL):
-    """Filter and function to get sine of the value."""
-    try:
-        return math.sin(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("sin", value)
-        return default
-
-
-def cosine(value, default=_SENTINEL):
-    """Filter and function to get cosine of the value."""
-    try:
-        return math.cos(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("cos", value)
-        return default
-
-
-def tangent(value, default=_SENTINEL):
-    """Filter and function to get tangent of the value."""
-    try:
-        return math.tan(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("tan", value)
-        return default
-
-
-def arc_sine(value, default=_SENTINEL):
-    """Filter and function to get arc sine of the value."""
-    try:
-        return math.asin(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("asin", value)
-        return default
-
-
-def arc_cosine(value, default=_SENTINEL):
-    """Filter and function to get arc cosine of the value."""
-    try:
-        return math.acos(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("acos", value)
-        return default
-
-
-def arc_tangent(value, default=_SENTINEL):
-    """Filter and function to get arc tangent of the value."""
-    try:
-        return math.atan(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("atan", value)
-        return default
-
-
-def arc_tangent2(*args, default=_SENTINEL):
-    """Filter and function to calculate four quadrant arc tangent of y / x.
-
-    The parameters to atan2 may be passed either in an iterable or as separate arguments
-    The default value may be passed either as a positional or in a keyword argument
-    """
-    try:
-        if 1 <= len(args) <= 2 and isinstance(args[0], (list, tuple)):
-            if len(args) == 2 and default is _SENTINEL:
-                # Default value passed as a positional argument
-                default = args[1]
-            args = args[0]
-        elif len(args) == 3 and default is _SENTINEL:
-            # Default value passed as a positional argument
-            default = args[2]
-
-        return math.atan2(float(args[0]), float(args[1]))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("atan2", args)
-        return default
-
-
 def version(value):
     """Filter and function to get version object of the value."""
     return AwesomeVersion(value)
 
 
-def square_root(value, default=_SENTINEL):
-    """Filter and function to get square root of the value."""
-    try:
-        return math.sqrt(float(value))
-    except (ValueError, TypeError):
-        if default is _SENTINEL:
-            raise_no_default("sqrt", value)
-        return default
-
-
 def timestamp_custom(value, date_format=DATE_STR_FORMAT, local=True, default=_SENTINEL):
     """Filter to convert given timestamp to format."""
     try:
@@ -2315,118 +2204,6 @@ def fail_when_undefined(value):
     return value
 
 
-def min_max_from_filter(builtin_filter: Any, name: str) -> Any:
-    """Convert a built-in min/max Jinja filter to a global function.
-
-    The parameters may be passed as an iterable or as separate arguments.
-    """
-
-    @pass_environment
-    @wraps(builtin_filter)
-    def wrapper(environment: jinja2.Environment, *args: Any, **kwargs: Any) -> Any:
-        if len(args) == 0:
-            raise TypeError(f"{name} expected at least 1 argument, got 0")
-
-        if len(args) == 1:
-            if isinstance(args[0], Iterable):
-                return builtin_filter(environment, args[0], **kwargs)
-
-            raise TypeError(f"'{type(args[0]).__name__}' object is not iterable")
-
-        return builtin_filter(environment, args, **kwargs)
-
-    return pass_environment(wrapper)
-
-
-def average(*args: Any, default: Any = _SENTINEL) -> Any:
-    """Filter and function to calculate the arithmetic mean.
-
-    Calculates of an iterable or of two or more arguments.
-
-    The parameters may be passed as an iterable or as separate arguments.
-    """
-    if len(args) == 0:
-        raise TypeError("average expected at least 1 argument, got 0")
-
-    # If first argument is iterable and more than 1 argument provided but not a named
-    # default, then use 2nd argument as default.
-    if isinstance(args[0], Iterable):
-        average_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:
-        average_list = args
-
-    try:
-        return statistics.fmean(average_list)
-    except (TypeError, statistics.StatisticsError):
-        if default is _SENTINEL:
-            raise_no_default("average", args)
-        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:
@@ -2549,21 +2326,6 @@ def regex_findall(value, find="", ignorecase=False):
     return _regex_cache(find, flags).findall(value)
 
 
-def bitwise_and(first_value, second_value):
-    """Perform a bitwise and operation."""
-    return first_value & second_value
-
-
-def bitwise_or(first_value, second_value):
-    """Perform a bitwise or operation."""
-    return first_value | second_value
-
-
-def bitwise_xor(first_value, second_value):
-    """Perform a bitwise xor operation."""
-    return first_value ^ second_value
-
-
 def struct_pack(value: Any | None, format_string: str) -> bytes | None:
     """Pack an object into a bytes object."""
     try:
@@ -3065,45 +2827,29 @@ def __init__(
         self.add_extension("jinja2.ext.do")
         self.add_extension("homeassistant.helpers.template.extensions.Base64Extension")
         self.add_extension("homeassistant.helpers.template.extensions.CryptoExtension")
+        self.add_extension("homeassistant.helpers.template.extensions.MathExtension")
 
-        self.globals["acos"] = arc_cosine
         self.globals["as_datetime"] = as_datetime
         self.globals["as_function"] = as_function
         self.globals["as_local"] = dt_util.as_local
         self.globals["as_timedelta"] = as_timedelta
         self.globals["as_timestamp"] = forgiving_as_timestamp
-        self.globals["asin"] = arc_sine
-        self.globals["atan"] = arc_tangent
-        self.globals["atan2"] = arc_tangent2
-        self.globals["average"] = average
         self.globals["bool"] = forgiving_boolean
         self.globals["combine"] = combine
-        self.globals["cos"] = cosine
         self.globals["difference"] = difference
-        self.globals["e"] = math.e
         self.globals["flatten"] = flatten
         self.globals["float"] = forgiving_float
         self.globals["iif"] = iif
         self.globals["int"] = forgiving_int
         self.globals["intersect"] = intersect
         self.globals["is_number"] = is_number
-        self.globals["log"] = logarithm
-        self.globals["max"] = min_max_from_filter(self.filters["max"], "max")
-        self.globals["median"] = median
         self.globals["merge_response"] = merge_response
-        self.globals["min"] = min_max_from_filter(self.filters["min"], "min")
         self.globals["pack"] = struct_pack
-        self.globals["pi"] = math.pi
         self.globals["set"] = _to_set
         self.globals["shuffle"] = shuffle
-        self.globals["sin"] = sine
         self.globals["slugify"] = slugify
-        self.globals["sqrt"] = square_root
-        self.globals["statistical_mode"] = statistical_mode
         self.globals["strptime"] = strptime
         self.globals["symmetric_difference"] = symmetric_difference
-        self.globals["tan"] = tangent
-        self.globals["tau"] = math.pi * 2
         self.globals["timedelta"] = timedelta
         self.globals["tuple"] = _to_tuple
         self.globals["typeof"] = typeof
@@ -3113,25 +2859,16 @@ def __init__(
         self.globals["version"] = version
         self.globals["zip"] = zip
 
-        self.filters["acos"] = arc_cosine
         self.filters["add"] = add
         self.filters["apply"] = apply
         self.filters["as_datetime"] = as_datetime
         self.filters["as_function"] = as_function
         self.filters["as_local"] = dt_util.as_local
         self.filters["as_timedelta"] = as_timedelta
         self.filters["as_timestamp"] = forgiving_as_timestamp
-        self.filters["asin"] = arc_sine
-        self.filters["atan"] = arc_tangent
-        self.filters["atan2"] = arc_tangent2
-        self.filters["average"] = average
-        self.filters["bitwise_and"] = bitwise_and
-        self.filters["bitwise_or"] = bitwise_or
-        self.filters["bitwise_xor"] = bitwise_xor
         self.filters["bool"] = forgiving_boolean
         self.filters["combine"] = combine
         self.filters["contains"] = contains
-        self.filters["cos"] = cosine
         self.filters["difference"] = difference
         self.filters["flatten"] = flatten
         self.filters["float"] = forgiving_float_filter
@@ -3142,8 +2879,6 @@ def __init__(
         self.filters["intersect"] = intersect
         self.filters["is_defined"] = fail_when_undefined
         self.filters["is_number"] = is_number
-        self.filters["log"] = logarithm
-        self.filters["median"] = median
         self.filters["multiply"] = multiply
         self.filters["ord"] = ord
         self.filters["ordinal"] = ordinal
@@ -3156,12 +2891,8 @@ def __init__(
         self.filters["regex_search"] = regex_search
         self.filters["round"] = forgiving_round
         self.filters["shuffle"] = shuffle
-        self.filters["sin"] = sine
         self.filters["slugify"] = slugify
-        self.filters["sqrt"] = square_root
-        self.filters["statistical_mode"] = statistical_mode
         self.filters["symmetric_difference"] = symmetric_difference
-        self.filters["tan"] = tangent
         self.filters["timestamp_custom"] = timestamp_custom
         self.filters["timestamp_local"] = timestamp_local
         self.filters["timestamp_utc"] = timestamp_utc

homeassistant/helpers/template/extensions/__init__.py

@@ -2,5 +2,6 @@
 
 from .base64 import Base64Extension
 from .crypto import CryptoExtension
+from .math import MathExtension
 
-__all__ = ["Base64Extension", "CryptoExtension"]
+__all__ = ["Base64Extension", "CryptoExtension", "MathExtension"]

homeassistant/helpers/template/extensions/base.py

@@ -19,7 +19,7 @@ class TemplateFunction:
     """Definition for a template function, filter, or test."""
 
     name: str
-    func: Callable[..., Any]
+    func: Callable[..., Any] | Any
     as_global: bool = False
     as_filter: bool = False
     as_test: bool = False