Make AST evaluate methods async

Also update `Function.__call__()` to async.

This would make it possible for `coalesce` to start async streams as
necessary during evaluation.

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

Author: shsms

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

@@ -35,7 +35,7 @@ def latest_sample(self) -> Sample[QuantityT] | None:
         return self._latest_sample
 
     @override
-    def evaluate(self) -> Sample[QuantityT] | None:
+    async def evaluate(self) -> Sample[QuantityT] | None:
         """Return the base value of the latest sample for this component."""
         if self._latest_sample is None:
             raise ValueError("Next value has not been fetched yet.")
@@ -71,9 +71,9 @@ class FunCall(AstNode[QuantityT]):
     function: Function[QuantityT]
 
     @override
-    def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
+    async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
         """Evaluate the function call with its arguments."""
-        return self.function()
+        return await self.function()
 
     @override
     def format(self, wrap: bool = False) -> str:
@@ -88,7 +88,7 @@ class Constant(AstNode[QuantityT]):
     value: QuantityT
 
     @override
-    def evaluate(self) -> QuantityT | None:
+    async def evaluate(self) -> QuantityT | None:
         """Return the constant value."""
         return self.value
 
@@ -106,10 +106,10 @@ class Add(AstNode[QuantityT]):
     right: AstNode[QuantityT]
 
     @override
-    def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
+    async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
         """Evaluate the addition of the left and right nodes."""
-        left = self.left.evaluate()
-        right = self.right.evaluate()
+        left = await self.left.evaluate()
+        right = await self.right.evaluate()
         match left, right:
             case Sample(), Sample():
                 if left.value is None:
@@ -161,10 +161,10 @@ class Sub(AstNode[QuantityT]):
     right: AstNode[QuantityT]
 
     @override
-    def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
+    async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
         """Evaluate the subtraction of the right node from the left node."""
-        left = self.left.evaluate()
-        right = self.right.evaluate()
+        left = await self.left.evaluate()
+        right = await self.right.evaluate()
         print("Sub.evaluate:", left, right)
         match left, right:
             case Sample(), Sample():
@@ -217,10 +217,10 @@ class Mul(AstNode[QuantityT]):
     right: AstNode[QuantityT]
 
     @override
-    def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
+    async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
         """Evaluate the multiplication of the left and right nodes."""
-        left = self.left.evaluate()
-        right = self.right.evaluate()
+        left = await self.left.evaluate()
+        right = await self.right.evaluate()
         match left, right:
             case Sample(), Sample():
                 if left.value is None:
@@ -271,10 +271,10 @@ class Div(AstNode[QuantityT]):
     right: AstNode[QuantityT]
 
     @override
-    def evaluate(self) -> QuantityT | None:
+    async def evaluate(self) -> QuantityT | None:
         """Evaluate the division of the left node by the right node."""
-        left = self.left.evaluate()
-        right = self.right.evaluate()
+        left = await self.left.evaluate()
+        right = await self.right.evaluate()
         match left, right:
             case Sample(), Sample():
                 if left.value is None:

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

@@ -20,7 +20,7 @@ class AstNode(abc.ABC, Generic[QuantityT]):
     span: tuple[int, int] | None = None
 
     @abc.abstractmethod
-    def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
+    async def evaluate(self) -> Sample[QuantityT] | QuantityT | None:
         """Evaluate the expression and return its numerical value."""
 
     @abc.abstractmethod

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

@@ -65,7 +65,7 @@ async def _run(self) -> None:
                     if comp.latest_sample is not None
                 )
 
-                res = self._root.evaluate()
+                res = await self._root.evaluate()
                 next_sample = res if isinstance(res, Sample) else Sample(timestamp, res)
                 await self._output_sender.send(next_sample)
             except (StopAsyncIteration, StopIteration):

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

@@ -29,7 +29,7 @@ def name(self) -> str:
         """Return the name of the function."""
 
     @abc.abstractmethod
-    def __call__(self) -> Sample[QuantityT] | QuantityT | None:
+    async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         """Call the function with the given arguments."""
 
     def format(self) -> str:
@@ -63,11 +63,11 @@ def name(self) -> str:
         return "COALESCE"
 
     @override
-    def __call__(self) -> Sample[QuantityT] | QuantityT | None:
+    async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         """Return the first non-None argument."""
         ts: datetime | None = None
         for param in self.params:
-            arg = param.evaluate()
+            arg = await param.evaluate()
             match arg:
                 case Sample(value=value, timestamp=timestamp):
                     if value is not None:
@@ -94,12 +94,12 @@ def name(self) -> str:
         return "MAX"
 
     @override
-    def __call__(self) -> Sample[QuantityT] | QuantityT | None:
+    async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         """Return the maximum of the arguments."""
         max_value: QuantityT | None = None
         ts: datetime | None = None
         for param in self.params:
-            arg = param.evaluate()
+            arg = await param.evaluate()
             match arg:
                 case Sample(value=value, timestamp=timestamp):
                     ts = timestamp
@@ -128,12 +128,12 @@ def name(self) -> str:
         return "MIN"
 
     @override
-    def __call__(self) -> Sample[QuantityT] | QuantityT | None:
+    async def __call__(self) -> Sample[QuantityT] | QuantityT | None:
         """Return the minimum of the arguments."""
         min_value: QuantityT | None = None
         ts: datetime | None = None
         for param in self.params:
-            arg = param.evaluate()
+            arg = await param.evaluate()
             match arg:
                 case Sample(value=value, timestamp=timestamp):
                     ts = timestamp