Update functions to fetch their own arguments

This would enable special forms like `coalesce` to fetch only what's
necessary.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/formulas/_ast.py

@@ -58,18 +58,16 @@ class FunCall(AstNode):
     """A function call in the formula."""
 
     function: Function
-    args: list[AstNode]
 
     @override
     def evaluate(self) -> float | None:
         """Evaluate the function call with its arguments."""
-        return self.function(arg.evaluate() for arg in self.args)
+        return self.function()
 
     @override
     def format(self, wrap: bool = False) -> str:
         """Return a string representation of the function call node."""
-        args_str = ", ".join(str(arg) for arg in self.args)
-        return f"{self.function.name}({args_str})"
+        return self.function.format()
 
 
 @dataclass(kw_only=True)

src/frequenz/sdk/timeseries/formulas/_formula.py

@@ -284,8 +284,7 @@ def coalesce(
                 right_nodes.append(_ast.Constant(value=item.base_value))
 
         new_root = _ast.FunCall(
-            function=Coalesce(),
-            args=[self.root] + right_nodes,
+            function=Coalesce([self.root] + right_nodes),
         )
 
         return FormulaBuilder(
@@ -317,8 +316,7 @@ def min(
                 right_nodes.append(_ast.Constant(value=item.base_value))
 
         new_root = _ast.FunCall(
-            function=Min(),
-            args=[self.root] + right_nodes,
+            function=Min([self.root] + right_nodes),
         )
 
         return FormulaBuilder(
@@ -350,8 +348,7 @@ def max(
                 right_nodes.append(_ast.Constant(value=item.base_value))
 
         new_root = _ast.FunCall(
-            function=Max(),
-            args=[self.root] + right_nodes,
+            function=Max([self.root] + right_nodes),
         )
 
         return FormulaBuilder(

src/frequenz/sdk/timeseries/formulas/_functions.py

@@ -6,33 +6,43 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Iterable
+from dataclasses import dataclass
 
 from typing_extensions import override
 
+from ._base_ast_node import AstNode
 
+
+@dataclass
 class Function(abc.ABC):
     """A function that can be called in a formula expression."""
 
+    params: list[AstNode]
+
     @property
     @abc.abstractmethod
     def name(self) -> str:
         """Return the name of the function."""
 
     @abc.abstractmethod
-    def __call__(self, args: Iterable[float | None]) -> float | None:
+    def __call__(self) -> float | None:
         """Call the function with the given arguments."""
 
+    def format(self) -> str:
+        """Return a string representation of the function."""
+        params_str = ", ".join(str(param) for param in self.params)
+        return f"{self.name}({params_str})"
+
     @classmethod
-    def from_string(cls, name: str) -> Function:
+    def from_string(cls, name: str, params: list[AstNode]) -> Function:
         """Create a function instance from its name."""
         match name.upper():
             case "COALESCE":
-                return Coalesce()
+                return Coalesce(params)
             case "MAX":
-                return Max()
+                return Max(params)
             case "MIN":
-                return Min()
+                return Min(params)
             case _:
                 raise ValueError(f"Unknown function name: {name}")
 
@@ -47,9 +57,10 @@ def name(self) -> str:
         return "COALESCE"
 
     @override
-    def __call__(self, args: Iterable[float | None]) -> float | None:
+    def __call__(self) -> float | None:
         """Return the first non-None argument."""
-        for arg in args:
+        for param in self.params:
+            arg = param.evaluate()
             if arg is not None:
                 return arg
         return None
@@ -65,10 +76,11 @@ def name(self) -> str:
         return "MAX"
 
     @override
-    def __call__(self, args: Iterable[float | None]) -> float | None:
+    def __call__(self) -> float | None:
         """Return the maximum of the arguments."""
         max_value: float | None = None
-        for arg in args:
+        for param in self.params:
+            arg = param.evaluate()
             if arg is None:
                 return None
             if max_value is None or arg > max_value:
@@ -86,10 +98,11 @@ def name(self) -> str:
         return "MIN"
 
     @override
-    def __call__(self, args: Iterable[float | None]) -> float | None:
+    def __call__(self) -> float | None:
         """Return the minimum of the arguments."""
         min_value: float | None = None
-        for arg in args:
+        for param in self.params:
+            arg = param.evaluate()
             if arg is None:
                 return None
             if min_value is None or arg < min_value:

src/frequenz/sdk/timeseries/formulas/_parser.py

@@ -143,7 +143,7 @@ def _parse_bracketed(self) -> AstNode | None:
 
     def _parse_function_call(self) -> AstNode | None:
         fn_name: _token.Token = next(self._lexer)
-        args: list[AstNode] = []
+        params: list[AstNode] = []
 
         token: _token.Token | None = self._lexer.peek()
         if token is None or not isinstance(token, _token.OpenParen):
@@ -153,12 +153,12 @@ def _parse_function_call(self) -> AstNode | None:
 
         _ = next(self._lexer)  # consume '('
         while True:
-            arg = self._parse_term()
-            if arg is None:
+            param = self._parse_term()
+            if param is None:
                 raise ValueError(
                     f"Expected argument in function call at position {fn_name.span}"
                 )
-            args.append(arg)
+            params.append(param)
 
             token = self._lexer.peek()
             if token is not None and isinstance(token, _token.Comma):
@@ -173,8 +173,7 @@ def _parse_function_call(self) -> AstNode | None:
 
         return _ast.FunCall(
             span=fn_name.span,
-            function=Function.from_string(fn_name.value),
-            args=args,
+            function=Function.from_string(fn_name.value, params),
         )
 
     def _parse_primary(self) -> AstNode | None: