Take custom constructors in FormulaEngine for creating quantity objects

This commit also updates the FormulaBuilders to take custom
constructors as arguments.

FormulaEngines were previously taking the type of the quantity, with
which, it could use only the default constructors to create objects.

Because support for using default constructors is going away,
FormulaEngines and FormulaBuilders need to know the exact constructor
to use for creating output quantities.

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

Author: shsms

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

@@ -12,6 +12,7 @@
 from datetime import datetime
 from math import isinf, isnan
 from typing import (
+    Callable,
     Dict,
     Generic,
     List,
@@ -63,23 +64,23 @@ def __init__(
         name: str,
         steps: List[FormulaStep],
         metric_fetchers: Dict[str, MetricFetcher[QuantityT]],
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
     ) -> None:
         """Create a `FormulaEngine` instance.
 
         Args:
             name: A name for the formula.
             steps: Steps for the engine to execute, in post-fix order.
             metric_fetchers: Fetchers for each metric stream the formula depends on.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
         """
         self._name = name
         self._steps = steps
         self._metric_fetchers: Dict[str, MetricFetcher[QuantityT]] = metric_fetchers
         self._first_run = True
-        self._output_type: type[QuantityT] = output_type
+        self._create_method: Callable[[float], QuantityT] = create_method
 
     async def _synchronize_metric_timestamps(
         self, metrics: Set[asyncio.Task[Optional[Sample[QuantityT]]]]
@@ -172,7 +173,7 @@ async def apply(self) -> Sample[QuantityT]:
         if isnan(res) or isinf(res):
             return Sample(metric_ts, None)
 
-        return Sample(metric_ts, self._output_type(res))
+        return Sample(metric_ts, self._create_method(res))
 
 
 _CompositionType = Union[
@@ -218,7 +219,7 @@ class _ComposableFormulaEngine(
 ):
     """A base class for formula engines."""
 
-    _output_type: Type[QuantityT]
+    _create_method: Callable[[float], QuantityT]
     _higher_order_builder: Type[_GenericHigherOrderBuilder]
     _task: asyncio.Task[None] | None = None
 
@@ -242,7 +243,7 @@ def __add__(
             A formula builder that can take further expressions, or can be built
                 into a formula engine.
         """
-        return self._higher_order_builder(self, self._output_type) + other  # type: ignore
+        return self._higher_order_builder(self, self._create_method) + other  # type: ignore
 
     def __sub__(
         self, other: _GenericEngine | _GenericHigherOrderBuilder
@@ -257,7 +258,7 @@ def __sub__(
             A formula builder that can take further expressions, or can be built
                 into a formula engine.
         """
-        return self._higher_order_builder(self, self._output_type) - other  # type: ignore
+        return self._higher_order_builder(self, self._create_method) - other  # type: ignore
 
     def __mul__(
         self, other: _GenericEngine | _GenericHigherOrderBuilder
@@ -272,7 +273,7 @@ def __mul__(
             A formula builder that can take further expressions, or can be built
                 into a formula engine.
         """
-        return self._higher_order_builder(self, self._output_type) * other  # type: ignore
+        return self._higher_order_builder(self, self._create_method) * other  # type: ignore
 
     def __truediv__(
         self, other: _GenericEngine | _GenericHigherOrderBuilder
@@ -287,7 +288,7 @@ def __truediv__(
             A formula builder that can take further expressions, or can be built
                 into a formula engine.
         """
-        return self._higher_order_builder(self, self._output_type) / other  # type: ignore
+        return self._higher_order_builder(self, self._create_method) / other  # type: ignore
 
 
 class FormulaEngine(
@@ -307,28 +308,28 @@ class FormulaEngine(
     def __init__(
         self,
         builder: FormulaBuilder[QuantityT],
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
     ) -> None:
         """Create a `FormulaEngine` instance.
 
         Args:
             builder: A `FormulaBuilder` instance to get the formula steps and metric
                 fetchers from.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
         """
         self._higher_order_builder = HigherOrderFormulaBuilder
         self._name: str = builder.name
         self._builder: FormulaBuilder[QuantityT] = builder
-        self._output_type: Type[QuantityT] = output_type  # type: ignore
+        self._create_method = create_method
         self._channel: Broadcast[Sample[QuantityT]] = Broadcast(self._name)
 
     async def _run(self) -> None:
         await self._builder.subscribe()
         steps, metric_fetchers = self._builder.finalize()
         evaluator = FormulaEvaluator[QuantityT](
-            self._name, steps, metric_fetchers, self._output_type
+            self._name, steps, metric_fetchers, self._create_method
         )
         sender = self._channel.new_sender()
         while True:
@@ -378,7 +379,7 @@ class FormulaEngine3Phase(
     def __init__(
         self,
         name: str,
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
         phase_streams: Tuple[
             FormulaEngine[QuantityT],
             FormulaEngine[QuantityT],
@@ -389,14 +390,14 @@ def __init__(
 
         Args:
             name: A name for the formula.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
             phase_streams: output streams of formula engines running per-phase formulas.
         """
         self._higher_order_builder = HigherOrderFormulaBuilder3Phase
         self._name: str = name
-        self._output_type = output_type
+        self._create_method = create_method
         self._channel: Broadcast[Sample3Phase[QuantityT]] = Broadcast(self._name)
         self._task: asyncio.Task[None] | None = None
         self._streams: tuple[
@@ -474,17 +475,17 @@ class FormulaBuilder(Generic[QuantityT]):
         add the values and return the result.
     """
 
-    def __init__(self, name: str, output_type: Type[QuantityT]) -> None:
+    def __init__(self, name: str, create_method: Callable[[float], QuantityT]) -> None:
         """Create a `FormulaBuilder` instance.
 
         Args:
             name: A name for the formula being built.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
         """
         self._name = name
-        self._output_type: Type[QuantityT] = output_type
+        self._create_method: Callable[[float], QuantityT] = create_method
         self._build_stack: List[FormulaStep] = []
         self._steps: List[FormulaStep] = []
         self._metric_fetchers: Dict[str, MetricFetcher[QuantityT]] = {}
@@ -648,7 +649,7 @@ def build(self) -> FormulaEngine[QuantityT]:
             A `FormulaEngine` instance.
         """
         self.finalize()
-        return FormulaEngine(self, output_type=self._output_type)
+        return FormulaEngine(self, create_method=self._create_method)
 
 
 class _BaseHOFormulaBuilder(ABC, Generic[QuantityT]):
@@ -657,16 +658,16 @@ class _BaseHOFormulaBuilder(ABC, Generic[QuantityT]):
     def __init__(
         self,
         engine: FormulaEngine[QuantityT] | FormulaEngine3Phase[QuantityT],
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
     ) -> None:
         """Create a `GenericHigherOrderFormulaBuilder` instance.
 
         Args:
             engine: A first input stream to create a builder with, so that python
                 operators `+, -, *, /` can be used directly on newly created instances.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
         """
         self._steps: deque[
             tuple[
@@ -675,7 +676,7 @@ def __init__(
             ]
         ] = deque()
         self._steps.append((TokenType.COMPONENT_METRIC, engine))
-        self._output_type: Type[QuantityT] = output_type
+        self._create_method: Callable[[float], QuantityT] = create_method
 
     @overload
     def _push(
@@ -854,7 +855,7 @@ def build(
         Returns:
             A `FormulaEngine` instance.
         """
-        builder = FormulaBuilder(name, self._output_type)
+        builder = FormulaBuilder(name, self._create_method)
         for typ, value in self._steps:
             if typ == TokenType.COMPONENT_METRIC:
                 assert isinstance(value, FormulaEngine)
@@ -888,9 +889,9 @@ def build(
             A `FormulaEngine3Phase` instance.
         """
         builders = [
-            FormulaBuilder(name, self._output_type),
-            FormulaBuilder(name, self._output_type),
-            FormulaBuilder(name, self._output_type),
+            FormulaBuilder(name, self._create_method),
+            FormulaBuilder(name, self._create_method),
+            FormulaBuilder(name, self._create_method),
         ]
         for typ, value in self._steps:
             if typ == TokenType.COMPONENT_METRIC:
@@ -909,7 +910,7 @@ def build(
                     builders[phase].push_oper(value)
         return FormulaEngine3Phase(
             name,
-            self._output_type,
+            self._create_method,
             (
                 builders[0].build(),
                 builders[1].build(),

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

@@ -10,7 +10,7 @@
 from collections import abc
 from dataclasses import dataclass
 from enum import Enum
-from typing import Generic, Type
+from typing import Callable, Generic
 
 from frequenz.channels import Sender
 
@@ -88,15 +88,15 @@ def _get_builder(
         self,
         name: str,
         component_metric_id: ComponentMetricId,
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
     ) -> ResampledFormulaBuilder[QuantityT]:
         builder = ResampledFormulaBuilder(
             self._namespace,
             name,
             self._channel_registry,
             self._resampler_subscription_sender,
             component_metric_id,
-            output_type,
+            create_method,
         )
         return builder
 

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

@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import Generic, Type
+from typing import Callable, Generic
 
 from frequenz.channels import Receiver, Sender
 
@@ -28,7 +28,7 @@ def __init__(  # pylint: disable=too-many-arguments
         channel_registry: ChannelRegistry,
         resampler_subscription_sender: Sender[ComponentMetricRequest],
         metric_id: ComponentMetricId,
-        output_type: Type[QuantityT],
+        create_method: Callable[[float], QuantityT],
     ) -> None:
         """Create a `ResampledFormulaBuilder` instance.
 
@@ -41,16 +41,16 @@ def __init__(  # pylint: disable=too-many-arguments
             resampler_subscription_sender: A sender to send metric requests to the
                 resampling actor.
             metric_id: A metric ID to fetch for all components in this formula.
-            output_type: A type object to generate the output `Sample` with.  If the
-                formula is for generating power values, this would be `Power`, for
-                example.
+            create_method: A method to generate the output `Sample` value with.  If the
+                formula is for generating power values, this would be
+                `Power.from_watts`, for example.
         """
         self._channel_registry = channel_registry
         self._resampler_subscription_sender = resampler_subscription_sender
         self._namespace = namespace
         self._metric_id = metric_id
         self._resampler_requests: list[ComponentMetricRequest] = []
-        super().__init__(formula_name, output_type)  # type: ignore[arg-type]
+        super().__init__(formula_name, create_method)  # type: ignore[arg-type]
 
     def _get_resampled_receiver(
         self, component_id: int, metric_id: ComponentMetricId

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -35,7 +35,7 @@ async def test_formula_composition(  # pylint: disable=too-many-locals
             logical_meter._namespace,  # pylint: disable=protected-access
             4,
             ComponentMetricId.ACTIVE_POWER,
-            Power,
+            Power.from_watts,
         )
         grid_power_recv = logical_meter.grid_power.new_receiver()
         battery_power_recv = battery_pool.power.new_receiver()

tests/timeseries/_formula_engine/utils.py

@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 from math import isclose
-from typing import Type
+from typing import Callable
 
 from frequenz.channels import Receiver
 
@@ -21,7 +21,7 @@ def get_resampled_stream(
     namespace: str,
     comp_id: int,
     metric_id: ComponentMetricId,
-    output_type: Type[QuantityT],
+    create_method: Callable[[float], QuantityT],
 ) -> Receiver[Sample[QuantityT]]:
     """Return the resampled data stream for the given component."""
     # Create a `FormulaBuilder` instance, just in order to reuse its
@@ -34,7 +34,7 @@ def get_resampled_stream(
         _data_pipeline._get()._channel_registry,
         _data_pipeline._get()._resampling_request_sender(),
         metric_id,
-        output_type,
+        create_method,
     )
     # Resampled data is always `Quantity` type, so we need to convert it to the desired
     # output type.
@@ -44,7 +44,7 @@ def get_resampled_stream(
     ).map(
         lambda sample: Sample(
             sample.timestamp,
-            None if sample.value is None else output_type(sample.value.base_value),
+            None if sample.value is None else create_method(sample.value.base_value),
         )
     )
     # pylint: enable=protected-access