WIP Quantity ABC

Author: llucax

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

@@ -5,12 +5,12 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable, cast
 
 from frequenz.channels import Sender
 
 from ...microgrid.component import ComponentMetricId
-from .._quantities import Current, Power, Quantity
+from .._quantities import Current, Power, Quantity, SupportsFloatT
 from ._formula_generators._formula_generator import (
     FormulaGenerator,
     FormulaGeneratorConfig,
@@ -59,14 +59,16 @@ def from_string(
         formula: str,
         component_metric_id: ComponentMetricId,
         *,
+        value_constructor: Callable[[float], SupportsFloatT],
         nones_are_zeros: bool = False,
-    ) -> FormulaEngine[Quantity]:
+    ) -> FormulaEngine[SupportsFloatT]:
         """Get a receiver for a manual formula.
 
         Args:
             formula: formula to execute.
             component_metric_id: The metric ID to use when fetching receivers from the
                 resampling actor.
+            value_constructor: A function to create new values with.
             nones_are_zeros: Whether to treat None values from the stream as 0s.  If
                 False, the returned value will be a None.
 
@@ -75,18 +77,22 @@ def from_string(
         """
         channel_key = formula + component_metric_id.value
         if channel_key in self._string_engines:
-            return self._string_engines[channel_key]
+            return cast(
+                FormulaEngine[SupportsFloatT], self._string_engines[channel_key]
+            )
 
-        builder = ResampledFormulaBuilder(
+        builder: ResampledFormulaBuilder[SupportsFloatT] = ResampledFormulaBuilder(
             self._namespace,
             formula,
             self._channel_registry,
             self._resampler_subscription_sender,
             component_metric_id,
-            Quantity,
+            value_constructor,
         )
         formula_engine = builder.from_string(formula, nones_are_zeros=nones_are_zeros)
-        self._string_engines[channel_key] = formula_engine
+        self._string_engines[channel_key] = cast(
+            FormulaEngine[Quantity], formula_engine
+        )
 
         return formula_engine
 

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

@@ -9,20 +9,19 @@
 from math import isinf, isnan
 from typing import Generic
 
-from .. import Sample
-from .._quantities import QuantityT
+from .._base_types import Sample, SupportsFloatT
 from ._formula_steps import FormulaStep, MetricFetcher
 
 
-class FormulaEvaluator(Generic[QuantityT]):
+class FormulaEvaluator(Generic[SupportsFloatT]):
     """A post-fix formula evaluator that operates on `Sample` receivers."""
 
     def __init__(
         self,
         name: str,
         steps: list[FormulaStep],
-        metric_fetchers: dict[str, MetricFetcher[QuantityT]],
-        create_method: Callable[[float], QuantityT],
+        metric_fetchers: dict[str, MetricFetcher[SupportsFloatT]],
+        create_method: Callable[[float], SupportsFloatT],
     ) -> None:
         """Create a `FormulaEngine` instance.
 
@@ -36,12 +35,14 @@ def __init__(
         """
         self._name = name
         self._steps = steps
-        self._metric_fetchers: dict[str, MetricFetcher[QuantityT]] = metric_fetchers
+        self._metric_fetchers: dict[str, MetricFetcher[SupportsFloatT]] = (
+            metric_fetchers
+        )
         self._first_run = True
-        self._create_method: Callable[[float], QuantityT] = create_method
+        self._create_method: Callable[[float], SupportsFloatT] = create_method
 
     async def _synchronize_metric_timestamps(
-        self, metrics: set[asyncio.Task[Sample[QuantityT] | None]]
+        self, metrics: set[asyncio.Task[Sample[SupportsFloatT] | None]]
     ) -> datetime:
         """Synchronize the metric streams.
 
@@ -88,7 +89,7 @@ async def _synchronize_metric_timestamps(
         self._first_run = False
         return latest_ts
 
-    async def apply(self) -> Sample[QuantityT]:
+    async def apply(self) -> Sample[SupportsFloatT]:
         """Fetch the latest metrics, apply the formula once and return the result.
 
         Returns:

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

@@ -16,7 +16,7 @@
 
 from ....microgrid import component, connection_manager
 from ....microgrid.component import ComponentMetricId
-from ..._quantities import QuantityT
+from ..._base_types import SupportsFloatT
 from .._formula_engine import FormulaEngine, FormulaEngine3Phase
 from .._resampled_formula_builder import ResampledFormulaBuilder
 
@@ -52,7 +52,7 @@ class FormulaGeneratorConfig:
     """The component IDs to use for generating the formula."""
 
 
-class FormulaGenerator(ABC, Generic[QuantityT]):
+class FormulaGenerator(ABC, Generic[SupportsFloatT]):
     """A class for generating formulas from the component graph."""
 
     def __init__(
@@ -83,8 +83,8 @@ def _get_builder(
         self,
         name: str,
         component_metric_id: ComponentMetricId,
-        create_method: Callable[[float], QuantityT],
-    ) -> ResampledFormulaBuilder[QuantityT]:
+        create_method: Callable[[float], SupportsFloatT],
+    ) -> ResampledFormulaBuilder[SupportsFloatT]:
         builder = ResampledFormulaBuilder(
             self._namespace,
             name,
@@ -140,5 +140,5 @@ def _get_grid_component_successors(self) -> set[component.Component]:
     @abstractmethod
     def generate(
         self,
-    ) -> FormulaEngine[QuantityT] | FormulaEngine3Phase[QuantityT]:
+    ) -> FormulaEngine[SupportsFloatT] | FormulaEngine3Phase[SupportsFloatT]:
         """Generate a formula engine, based on the component graph."""

src/frequenz/sdk/timeseries/formula_engine/_formula_generators/_formula_generator.py

@@ -7,12 +7,14 @@
 
 import math
 from abc import ABC, abstractmethod
-from typing import Generic
+from typing import Generic, SupportsFloat, TypeVar
 
 from frequenz.channels import Receiver
 
 from .. import Sample
-from .._quantities import QuantityT
+
+SupportsFloatT = TypeVar("SupportsFloatT", bound=SupportsFloat)
+"""Type variable for types that support conversion to float."""
 
 
 class FormulaStep(ABC):
@@ -343,13 +345,13 @@ def apply(self, eval_stack: list[float]) -> None:
         eval_stack.append(val)
 
 
-class MetricFetcher(Generic[QuantityT], FormulaStep):
+class MetricFetcher(Generic[SupportsFloatT], FormulaStep):
     """A formula step for fetching a value from a metric Receiver."""
 
     def __init__(
         self,
         name: str,
-        stream: Receiver[Sample[QuantityT]],
+        stream: Receiver[Sample[SupportsFloatT]],
         *,
         nones_are_zeros: bool,
     ) -> None:
@@ -361,12 +363,12 @@ def __init__(
             nones_are_zeros: Whether to treat None values from the stream as 0s.
         """
         self._name = name
-        self._stream: Receiver[Sample[QuantityT]] = stream
-        self._next_value: Sample[QuantityT] | None = None
+        self._stream: Receiver[Sample[SupportsFloatT]] = stream
+        self._next_value: Sample[SupportsFloatT] | None = None
         self._nones_are_zeros = nones_are_zeros
 
     @property
-    def stream(self) -> Receiver[Sample[QuantityT]]:
+    def stream(self) -> Receiver[Sample[SupportsFloatT]]:
         """Return the stream from which to fetch values.
 
         Returns:
@@ -382,7 +384,7 @@ def stream_name(self) -> str:
         """
         return str(self._stream.__doc__)
 
-    async def fetch_next(self) -> Sample[QuantityT] | None:
+    async def fetch_next(self) -> Sample[SupportsFloatT] | None:
         """Fetch the next value from the stream.
 
         To be called before each call to `apply`.
@@ -394,7 +396,7 @@ async def fetch_next(self) -> Sample[QuantityT] | None:
         return self._next_value
 
     @property
-    def value(self) -> Sample[QuantityT] | None:
+    def value(self) -> Sample[SupportsFloatT] | None:
         """Get the next value in the stream.
 
         Returns:
@@ -423,10 +425,10 @@ def apply(self, eval_stack: list[float]) -> None:
             raise RuntimeError("No next value available to append.")
 
         next_value = self._next_value.value
-        if next_value is None or next_value.isnan() or next_value.isinf():
+        if next_value is None or math.isnan(next_value) or math.isinf(next_value):
             if self._nones_are_zeros:
                 eval_stack.append(0.0)
             else:
                 eval_stack.append(math.nan)
         else:
-            eval_stack.append(next_value.base_value)
+            eval_stack.append(float(next_value))

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

@@ -11,8 +11,7 @@
 from frequenz.channels import Receiver, Sender
 
 from ...microgrid.component import ComponentMetricId
-from .. import Sample
-from .._quantities import Quantity, QuantityT
+from .._base_types import Sample, SupportsFloatT
 from ._formula_engine import FormulaBuilder, FormulaEngine
 from ._tokenizer import Tokenizer, TokenType
 
@@ -21,7 +20,7 @@
     from ...actor import ChannelRegistry, ComponentMetricRequest
 
 
-class ResampledFormulaBuilder(FormulaBuilder[QuantityT]):
+class ResampledFormulaBuilder(FormulaBuilder[SupportsFloatT]):
     """Provides a way to build a FormulaEngine from resampled data streams."""
 
     def __init__(  # pylint: disable=too-many-arguments
@@ -31,7 +30,7 @@ def __init__(  # pylint: disable=too-many-arguments
         channel_registry: ChannelRegistry,
         resampler_subscription_sender: Sender[ComponentMetricRequest],
         metric_id: ComponentMetricId,
-        create_method: Callable[[float], QuantityT],
+        create_method: Callable[[float], SupportsFloatT],
     ) -> None:
         """Create a `ResampledFormulaBuilder` instance.
 
@@ -55,11 +54,11 @@ def __init__(  # pylint: disable=too-many-arguments
         self._namespace: str = namespace
         self._metric_id: ComponentMetricId = metric_id
         self._resampler_requests: list[ComponentMetricRequest] = []
-        super().__init__(formula_name, create_method)  # type: ignore[arg-type]
+        super().__init__(formula_name, create_method)
 
     def _get_resampled_receiver(
         self, component_id: int, metric_id: ComponentMetricId
-    ) -> Receiver[Sample[QuantityT]]:
+    ) -> Receiver[Sample[SupportsFloatT]]:
         """Get a receiver with the resampled data for the given component id.
 
         Args:
@@ -75,13 +74,13 @@ def _get_resampled_receiver(
         request = ComponentMetricRequest(self._namespace, component_id, metric_id, None)
         self._resampler_requests.append(request)
         resampled_channel = self._channel_registry.get_or_create(
-            Sample[Quantity], request.get_channel_name()
+            Sample[SupportsFloatT], request.get_channel_name()
         )
         resampled_receiver = resampled_channel.new_receiver().map(
             lambda sample: Sample(
                 sample.timestamp,
                 (
-                    self._create_method(sample.value.base_value)
+                    self._create_method(float(sample.value))
                     if sample.value is not None
                     else None
                 ),
@@ -112,7 +111,7 @@ def from_string(
         formula: str,
         *,
         nones_are_zeros: bool,
-    ) -> FormulaEngine[QuantityT]:
+    ) -> FormulaEngine[SupportsFloatT]:
         """Construct a `FormulaEngine` from the given formula string.
 
         Formulas can have Component IDs that are preceeded by a pound symbol("#"), and

src/frequenz/sdk/timeseries/formula_engine/_resampled_formula_builder.py

@@ -5,12 +5,13 @@
 
 
 import uuid
+from typing import Callable
 
 from frequenz.channels import Sender
 
 from ...actor import ChannelRegistry, ComponentMetricRequest
 from ...microgrid.component import ComponentMetricId
-from .._quantities import Power, Quantity
+from .._quantities import Power, SupportsFloatT
 from ..formula_engine import FormulaEngine
 from ..formula_engine._formula_engine_pool import FormulaEnginePool
 from ..formula_engine._formula_generators import CHPPowerFormula, PVPowerFormula
@@ -99,8 +100,9 @@ def start_formula(
         formula: str,
         component_metric_id: ComponentMetricId,
         *,
+        value_constructor: Callable[[float], SupportsFloatT],
         nones_are_zeros: bool = False,
-    ) -> FormulaEngine[Quantity]:
+    ) -> FormulaEngine[SupportsFloatT]:
         """Start execution of the given formula.
 
         Formulas can have Component IDs that are preceeded by a pound symbol("#"), and
@@ -113,14 +115,18 @@ def start_formula(
             formula: formula to execute.
             component_metric_id: The metric ID to use when fetching receivers from the
                 resampling actor.
+            value_constructor: A function to create new values with.
             nones_are_zeros: Whether to treat None values from the stream as 0s.  If
                 False, the returned value will be a None.
 
         Returns:
             A FormulaEngine that applies the formula and streams values.
         """
         return self._formula_pool.from_string(
-            formula, component_metric_id, nones_are_zeros=nones_are_zeros
+            formula,
+            component_metric_id,
+            value_constructor=value_constructor,
+            nones_are_zeros=nones_are_zeros,
         )
 
     @property

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

@@ -12,7 +12,7 @@
 from frequenz.sdk.microgrid import _data_pipeline
 from frequenz.sdk.microgrid.component import ComponentMetricId
 from frequenz.sdk.timeseries import Sample
-from frequenz.sdk.timeseries._quantities import QuantityT
+from frequenz.sdk.timeseries._quantities import SupportsFloatT
 from frequenz.sdk.timeseries.formula_engine._resampled_formula_builder import (
     ResampledFormulaBuilder,
 )
@@ -22,8 +22,8 @@ def get_resampled_stream(
     namespace: str,
     comp_id: int,
     metric_id: ComponentMetricId,
-    create_method: Callable[[float], QuantityT],
-) -> Receiver[Sample[QuantityT]]:
+    create_method: Callable[[float], SupportsFloatT],
+) -> Receiver[Sample[SupportsFloatT]]:
     """Return the resampled data stream for the given component."""
     # Create a `FormulaBuilder` instance, just in order to reuse its
     # `_get_resampled_receiver` function implementation.
@@ -51,10 +51,10 @@ def get_resampled_stream(
     # pylint: enable=protected-access
 
 
-def equal_float_lists(list1: list[QuantityT], list2: list[QuantityT]) -> bool:
+def equal_float_lists(list1: list[SupportsFloatT], list2: list[SupportsFloatT]) -> bool:
     """Compare two float lists with `math.isclose()`."""
     return (
         len(list1) > 0
         and len(list1) == len(list2)
-        and all(isclose(v1.base_value, v2.base_value) for v1, v2 in zip(list1, list2))
+        and all(isclose(v1, v2) for v1, v2 in zip(list1, list2))
     )