Add support for calculating averages in FormulaEngine

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

Author: shsms

src/frequenz/sdk/timeseries/logical_meter/_exceptions.py

@@ -0,0 +1,8 @@
+# License: MIT
+# Copyright © 2022 Frequenz Energy-as-a-Service GmbH
+
+"""Formula Engine Exceptions."""
+
+
+class FormulaEngineError(Exception):
+    """An error occured while fetching metrics or applying the formula on them."""

src/frequenz/sdk/timeseries/logical_meter/_formula_engine.py

@@ -7,13 +7,14 @@
 
 import asyncio
 from datetime import datetime
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, Tuple
 
 from frequenz.channels import Receiver
 
 from .. import Sample
 from ._formula_steps import (
     Adder,
+    Averager,
     Divider,
     FormulaStep,
     MetricFetcher,
@@ -218,6 +219,20 @@ def push_metric(
         )
         self._steps.append(fetcher)
 
+    def push_average(self, metrics: List[Tuple[str, Receiver[Sample], bool]]) -> None:
+        """Push an average calculator into the engine.
+
+        Args:
+            metrics: list of arguments to pass to each `MetricFetcher`.
+        """
+        fetchers: List[MetricFetcher] = []
+        for metric in metrics:
+            fetcher = self._metric_fetchers.setdefault(
+                metric[0], MetricFetcher(*metric)
+            )
+            fetchers.append(fetcher)
+        self._steps.append(Averager(fetchers))
+
     def build(self) -> FormulaEngine:
         """Finalize and build the formula engine.
 

src/frequenz/sdk/timeseries/logical_meter/_formula_steps.py

@@ -3,12 +3,15 @@
 
 """Steps for building formula engines with."""
 
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import List, Optional
 
 from frequenz.channels import Receiver
 
 from .. import Sample
+from ._exceptions import FormulaEngineError
 
 
 class FormulaStep(ABC):
@@ -157,6 +160,54 @@ def apply(self, _: List[Optional[float]]) -> None:
         """No-op."""
 
 
+class Averager(FormulaStep):
+    """A formula step for calculating average."""
+
+    def __init__(self, fetchers: List[MetricFetcher]) -> None:
+        """Create an `Averager` instance.
+
+        Args:
+            fetchers: MetricFetchers for the metrics that need to be averaged.
+        """
+        self._fetchers = fetchers
+
+    def __repr__(self) -> str:
+        """Return a string representation of the step.
+
+        Returns:
+            A string representation of the step.
+        """
+        return f"avg({', '.join(repr(f) for f in self._fetchers)})"
+
+    def apply(self, eval_stack: List[Optional[float]]) -> None:
+        """Calculate average of given metrics, push the average to the eval_stack.
+
+        Args:
+            eval_stack: An evaluation stack, to append the calculated average to.
+
+        Raises:
+            FormulaEngineError: when metric fetchers are unable to fetch values.
+        """
+        value_count = 0
+        total = 0.0
+        for fetcher in self._fetchers:
+            next_val = fetcher.value
+            if next_val is None:
+                raise FormulaEngineError(
+                    "Unable to fetch a value from the resampling actor."
+                )
+            if next_val.value is None:
+                continue
+            value_count += 1
+            total += next_val.value
+        if value_count == 0:
+            avg = 0.0
+        else:
+            avg = total / value_count
+
+        eval_stack.append(avg)
+
+
 class MetricFetcher(FormulaStep):
     """A formula step for fetching a value from a metric Receiver."""
 
@@ -186,6 +237,15 @@ async def fetch_next(self) -> Optional[Sample]:
         self._next_value = await self._stream.receive()
         return self._next_value
 
+    @property
+    def value(self) -> Optional[Sample]:
+        """Get the next value in the stream.
+
+        Returns:
+            Next value in the stream.
+        """
+        return self._next_value
+
     def __repr__(self) -> str:
         """Return a string representation of the step.