Update remaining code to work with FormulaEngine interface changes

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

Author: shsms

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

@@ -2,24 +2,13 @@
 # Copyright © 2023 Frequenz Energy-as-a-Service GmbH
 
 """A formula engine for applying formulas."""
-from ._formula_engine import (
-    FormulaEngine,
-    FormulaEngine3Phase,
-    FormulaReceiver,
-    FormulaReceiver3Phase,
-    _GenericEngine,
-    _GenericFormulaReceiver,
-)
+from ._formula_engine import FormulaEngine, FormulaEngine3Phase
 from ._formula_engine_pool import FormulaEnginePool
 from ._resampled_formula_builder import ResampledFormulaBuilder
 
 __all__ = [
     "FormulaEngine",
     "FormulaEngine3Phase",
-    "FormulaReceiver",
-    "FormulaReceiver3Phase",
     "FormulaEnginePool",
-    "_GenericEngine",
-    "_GenericFormulaReceiver",
     "ResampledFormulaBuilder",
 ]

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

@@ -19,12 +19,7 @@
 
 if TYPE_CHECKING:
     # Break circular import by enclosing these type hints in a `TYPE_CHECKING` block.
-    from .._formula_engine import (
-        FormulaReceiver,
-        FormulaReceiver3Phase,
-        _GenericEngine,
-        _GenericFormulaReceiver,
-    )
+    from .._formula_engine import FormulaEngine, FormulaEngine3Phase
 
 
 class FormulaEnginePool:
@@ -51,14 +46,14 @@ def __init__(
         self._namespace = namespace
         self._channel_registry = channel_registry
         self._resampler_subscription_sender = resampler_subscription_sender
-        self._engines: dict[str, "FormulaReceiver|FormulaReceiver3Phase"] = {}
+        self._engines: dict[str, "FormulaEngine|FormulaEngine3Phase"] = {}
 
-    async def from_string(
+    def from_string(
         self,
         formula: str,
         component_metric_id: ComponentMetricId,
         nones_are_zeros: bool = False,
-    ) -> "FormulaReceiver":
+    ) -> FormulaEngine:
         """Get a receiver for a manual formula.
 
         Args:
@@ -73,7 +68,7 @@ async def from_string(
         """
         channel_key = formula + component_metric_id.value
         if channel_key in self._engines:
-            return self._engines[channel_key].new_receiver()
+            return self._engines[channel_key]  # type: ignore
 
         builder = ResampledFormulaBuilder(
             self._namespace,
@@ -82,17 +77,17 @@ async def from_string(
             self._resampler_subscription_sender,
             component_metric_id,
         )
-        formula_engine = await builder.from_string(formula, nones_are_zeros)
+        formula_engine = builder.from_string(formula, nones_are_zeros)
         self._engines[channel_key] = formula_engine
 
-        return formula_engine.new_receiver()
+        return formula_engine
 
-    async def from_generator(
+    def from_generator(
         self,
         channel_key: str,
-        generator: "Type[FormulaGenerator[_GenericEngine]]",
+        generator: "Type[FormulaGenerator]",
         config: FormulaGeneratorConfig = FormulaGeneratorConfig(),
-    ) -> "_GenericFormulaReceiver":
+    ) -> FormulaEngine | FormulaEngine3Phase:
         """Get a receiver for a formula from a generator.
 
         Args:
@@ -105,13 +100,13 @@ async def from_generator(
                 FormulaGenerator returns.
         """
         if channel_key in self._engines:
-            return self._engines[channel_key].new_receiver()
+            return self._engines[channel_key]
 
-        engine = await generator(
+        engine = generator(
             self._namespace,
             self._channel_registry,
             self._resampler_subscription_sender,
             config,
         ).generate()
         self._engines[channel_key] = engine
-        return engine.new_receiver()
+        return engine

src/frequenz/sdk/timeseries/ev_charger_pool/_ev_charger_pool.py

@@ -19,7 +19,7 @@
 from ...microgrid import connection_manager
 from ...microgrid.component import ComponentCategory, ComponentMetricId
 from .. import Sample, Sample3Phase
-from .._formula_engine import FormulaEnginePool, FormulaReceiver, FormulaReceiver3Phase
+from .._formula_engine import FormulaEngine, FormulaEngine3Phase, FormulaEnginePool
 from .._formula_engine._formula_generators import (
     EVChargerCurrentFormula,
     EVChargerPowerFormula,
@@ -104,40 +104,45 @@ def component_ids(self) -> abc.Set[int]:
         """
         return self._component_ids
 
-    async def current(self) -> FormulaReceiver3Phase:
+    def current(self) -> FormulaEngine3Phase:
         """Fetch the total current for the EV Chargers in the pool.
 
         If a formula engine to calculate EV Charger current is not already running, it
-        will be started.  Else, it will return a new receiver to the already existing
-        data stream.
+        will be started.
+
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
         Returns:
-            A *new* receiver that will stream EV Charger current values.
+            A FormulaEngine that will calculate and stream the total current of all EV
+                Chargers.
         """
-        return await self._formula_pool.from_generator(
+        return self._formula_pool.from_generator(
             "ev_charger_total_current",
             EVChargerCurrentFormula,
             FormulaGeneratorConfig(component_ids=self._component_ids),
-        )
+        )  # type: ignore[return-value]
 
-    async def power(self) -> FormulaReceiver:
+    def power(self) -> FormulaEngine:
         """Fetch the total power for the EV Chargers in the pool.
 
         If a formula engine to calculate EV Charger power is not already running, it
-        will be started.  Else, it will return a new receiver to the already existing
-        data stream.
+        will be started.
 
-        Returns:
-            A *new* receiver that will stream EV Charger power values.
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
+        Returns:
+            A FormulaEngine that will calculate and stream the total power of all EV
+                Chargers.
         """
-        return await self._formula_pool.from_generator(
+        return self._formula_pool.from_generator(
             "ev_charger_total_power",
             EVChargerPowerFormula,
             FormulaGeneratorConfig(component_ids=self._component_ids),
-        )
+        )  # type: ignore[return-value]
 
-    async def component_data(self, component_id: int) -> Receiver[EVChargerData]:
+    def component_data(self, component_id: int) -> Receiver[EVChargerData]:
         """Stream 3-phase current values and state of an EV Charger.
 
         Args:

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

@@ -13,7 +13,7 @@
 from ...actor import ChannelRegistry, ComponentMetricRequest
 from ...microgrid import ComponentGraph
 from ...microgrid.component import ComponentMetricId
-from .._formula_engine import FormulaEnginePool, FormulaReceiver, FormulaReceiver3Phase
+from .._formula_engine import FormulaEngine, FormulaEngine3Phase, FormulaEnginePool
 from .._formula_engine._formula_generators import (
     BatteryPowerFormula,
     GridCurrentFormula,
@@ -112,12 +112,12 @@ def __init__(
             self._resampler_subscription_sender,
         )
 
-    async def start_formula(
+    def start_formula(
         self,
         formula: str,
         component_metric_id: ComponentMetricId,
         nones_are_zeros: bool = False,
-    ) -> FormulaReceiver:
+    ) -> FormulaEngine:
         """Start execution of the given formula.
 
         Formulas can have Component IDs that are preceeded by a pound symbol("#"), and
@@ -134,64 +134,76 @@ async def start_formula(
                 False, the returned value will be a None.
 
         Returns:
-            A FormulaReceiver that streams values with the formulas applied.
+            A FormulaEngine that applies the formula and streams values.
         """
-        return await self._formula_pool.from_string(
+        return self._formula_pool.from_string(
             formula, component_metric_id, nones_are_zeros
         )
 
-    async def grid_power(self) -> FormulaReceiver:
+    def grid_power(self) -> FormulaEngine:
         """Fetch the grid power for the microgrid.
 
         If a formula engine to calculate grid power is not already running, it will be
-        started.  Else, it will return a new receiver to the already existing data
-        stream.
+        started.
 
-        Returns:
-            A *new* receiver that will stream grid_power values.
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
+        Returns:
+            A FormulaEngine that will calculate and stream grid power.
         """
-        return await self._formula_pool.from_generator("grid_power", GridPowerFormula)
+        return self._formula_pool.from_generator(
+            "grid_power",
+            GridPowerFormula,
+        )  # type: ignore[return-value]
 
-    async def grid_current(self) -> FormulaReceiver3Phase:
+    def grid_current(self) -> FormulaEngine3Phase:
         """Fetch the grid power for the microgrid.
 
         If a formula engine to calculate grid current is not already running, it will be
-        started.  Else, it will return a new receiver to the already existing data
-        stream.
+        started.
 
-        Returns:
-            A *new* receiver that will stream grid_current values.
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
+        Returns:
+            A FormulaEngine that will calculate and stream grid current.
         """
-        return await self._formula_pool.from_generator(
-            "grid_current", GridCurrentFormula
-        )
+        return self._formula_pool.from_generator(
+            "grid_current",
+            GridCurrentFormula,
+        )  # type: ignore[return-value]
 
-    async def battery_power(self) -> FormulaReceiver:
+    def battery_power(self) -> FormulaEngine:
         """Fetch the cumulative battery power in the microgrid.
 
         If a formula engine to calculate cumulative battery power is not already
-        running, it will be started.  Else, it will return a new receiver to the already
-        existing data stream.
+        running, it will be started.
 
-        Returns:
-            A *new* receiver that will stream battery_power values.
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
+        Returns:
+            A FormulaEngine that will calculate and stream battery power.
         """
-        return await self._formula_pool.from_generator(
-            "battery_power", BatteryPowerFormula
-        )
+        return self._formula_pool.from_generator(
+            "battery_power",
+            BatteryPowerFormula,
+        )  # type: ignore[return-value]
 
-    async def pv_power(self) -> FormulaReceiver:
+    def pv_power(self) -> FormulaEngine:
         """Fetch the PV power production in the microgrid.
 
         If a formula engine to calculate PV power production is not already running, it
-        will be started.  Else, it will return a new receiver to the already existing
-        data stream.
+        will be started.
 
-        Returns:
-            A *new* receiver that will stream PV power production values.
+        A receiver from the formula engine can be created using the `new_receiver`
+        method.
 
+        Returns:
+            A FormulaEngine that will calculate and stream PV power production.
         """
-        return await self._formula_pool.from_generator("pv_power", PVPowerFormula)
+        return self._formula_pool.from_generator(
+            "pv_power",
+            PVPowerFormula,
+        )  # type: ignore[return-value]