Make Samples generic over Quantity types

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

Author: shsms

benchmarks/timeseries/periodic_feature_extractor.py

@@ -30,7 +30,7 @@
 async def init_feature_extractor(period: int) -> PeriodicFeatureExtractor:
     """Initialize the PeriodicFeatureExtractor class."""
     # We only need the moving window to initialize the PeriodicFeatureExtractor class.
-    lm_chan = Broadcast[Sample]("lm_net_power")
+    lm_chan = Broadcast[Sample[Quantity]]("lm_net_power")
     moving_window = MovingWindow(
         timedelta(seconds=1), lm_chan.new_receiver(), timedelta(seconds=1)
     )

benchmarks/timeseries/resampling.py

@@ -17,7 +17,7 @@
 
 
 def nop(  # pylint: disable=unused-argument
-    samples: Sequence[Sample],
+    samples: Sequence[Sample[Quantity]],
     resampler_config: ResamplerConfig,
     source_properties: SourceProperties,
 ) -> float:

examples/power_distribution.py

@@ -22,6 +22,7 @@
 from frequenz.sdk.actor import ResamplerConfig
 from frequenz.sdk.actor.power_distributing import Result, Success
 from frequenz.sdk.timeseries import Sample
+from frequenz.sdk.timeseries._quantities import Quantity
 
 _logger = logging.getLogger(__name__)
 HOST = "microgrid.sandbox.api.frequenz.io"  # it should be the host name.
@@ -97,7 +98,7 @@ class DataCollectingActor:
     def __init__(
         self,
         request_channel: Sender[List[float]],
-        active_power_data: Receiver[Sample],
+        active_power_data: Receiver[Sample[Quantity]],
     ) -> None:
         """Create actor instance.
 

examples/resampling.py

@@ -38,7 +38,7 @@ async def _calculate_average(source: Source, sink: Sink) -> None:
         await sink(Sample(datetime.now(timezone.utc), Quantity(avg)))
 
 
-async def _print_sample(sample: Sample) -> None:
+async def _print_sample(sample: Sample[Quantity]) -> None:
     print(f"\nResampled average at {sample.timestamp}: {sample.value}\n")
 
 
@@ -106,7 +106,7 @@ async def run() -> None:  # pylint: disable=too-many-locals
     )
 
     # Create a channel to calculate an average for all the data
-    average_chan = Broadcast[Sample]("average")
+    average_chan = Broadcast[Sample[Quantity]]("average")
 
     second_stage_resampler = Resampler(
         ResamplerConfig(resampling_period=timedelta(seconds=3.0))
@@ -118,7 +118,7 @@ async def run() -> None:  # pylint: disable=too-many-locals
     average_sender = average_chan.new_sender()
 
     # Needed until channels Senders raises exceptions on errors
-    async def sink_adapter(sample: Sample) -> None:
+    async def sink_adapter(sample: Sample[Quantity]) -> None:
         await average_sender.send(sample)
 
     print("Starting...")

src/frequenz/sdk/actor/_data_sourcing/microgrid_api_source.py

@@ -305,7 +305,7 @@ def _get_metric_senders(
         self,
         category: ComponentCategory,
         requests: Dict[ComponentMetricId, List[ComponentMetricRequest]],
-    ) -> List[Tuple[Callable[[Any], float], List[Sender[Sample]]]]:
+    ) -> List[Tuple[Callable[[Any], float], List[Sender[Sample[Quantity]]]]]:
         """Get channel senders from the channel registry for each requested metric.
 
         Args:

src/frequenz/sdk/actor/_resampling.py

@@ -13,6 +13,7 @@
 
 from .._internal._asyncio import cancel_and_await
 from ..timeseries import Sample
+from ..timeseries._quantities import Quantity
 from ..timeseries._resampling import Resampler, ResamplerConfig, ResamplingError
 from ._channel_registry import ChannelRegistry
 from ._data_sourcing import ComponentMetricRequest
@@ -80,7 +81,7 @@ async def _subscribe(self, request: ComponentMetricRequest) -> None:
         # exceptions to report errors.
         sender = self._channel_registry.new_sender(request.get_channel_name())
 
-        async def sink_adapter(sample: Sample) -> None:
+        async def sink_adapter(sample: Sample[Quantity]) -> None:
             await sender.send(sample)
 
         self._resampler.add_timeseries(request_channel_name, receiver, sink_adapter)

src/frequenz/sdk/timeseries/_base_types.py

@@ -6,16 +6,16 @@
 import functools
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import Callable, Iterator, Self, overload
+from typing import Callable, Generic, Iterator, Self, overload
 
-from ._quantities import Quantity
+from ._quantities import QuantityT
 
 UNIX_EPOCH = datetime.fromtimestamp(0.0, tz=timezone.utc)
 """The UNIX epoch (in UTC)."""
 
 
 @dataclass(frozen=True, order=True)
-class Sample:
+class Sample(Generic[QuantityT]):
     """A measurement taken at a particular point in time.
 
     The `value` could be `None` if a component is malfunctioning or data is
@@ -26,12 +26,12 @@ class Sample:
     timestamp: datetime
     """The time when this sample was generated."""
 
-    value: Quantity | None = None
+    value: QuantityT | None = None
     """The value of this sample."""
 
 
 @dataclass(frozen=True)
-class Sample3Phase:
+class Sample3Phase(Generic[QuantityT]):
     """A 3-phase measurement made at a particular point in time.
 
     Each of the `value` fields could be `None` if a component is malfunctioning
@@ -42,16 +42,16 @@ class Sample3Phase:
 
     timestamp: datetime
     """The time when this sample was generated."""
-    value_p1: Quantity | None
+    value_p1: QuantityT | None
     """The value of the 1st phase in this sample."""
 
-    value_p2: Quantity | None
+    value_p2: QuantityT | None
     """The value of the 2nd phase in this sample."""
 
-    value_p3: Quantity | None
+    value_p3: QuantityT | None
     """The value of the 3rd phase in this sample."""
 
-    def __iter__(self) -> Iterator[Quantity | None]:
+    def __iter__(self) -> Iterator[QuantityT | None]:
         """Return an iterator that yields values from each of the phases.
 
         Yields:
@@ -62,14 +62,14 @@ def __iter__(self) -> Iterator[Quantity | None]:
         yield self.value_p3
 
     @overload
-    def max(self, default: Quantity) -> Quantity:
+    def max(self, default: QuantityT) -> QuantityT:
         ...
 
     @overload
-    def max(self, default: None = None) -> Quantity | None:
+    def max(self, default: None = None) -> QuantityT | None:
         ...
 
-    def max(self, default: Quantity | None = None) -> Quantity | None:
+    def max(self, default: QuantityT | None = None) -> QuantityT | None:
         """Return the max value among all phases, or default if they are all `None`.
 
         Args:
@@ -80,21 +80,21 @@ def max(self, default: Quantity | None = None) -> Quantity | None:
         """
         if not any(self):
             return default
-        value: Quantity = functools.reduce(
+        value: QuantityT = functools.reduce(
             lambda x, y: x if x > y else y,
             filter(None, self),
         )
         return value
 
     @overload
-    def min(self, default: Quantity) -> Quantity:
+    def min(self, default: QuantityT) -> QuantityT:
         ...
 
     @overload
-    def min(self, default: None = None) -> Quantity | None:
+    def min(self, default: None = None) -> QuantityT | None:
         ...
 
-    def min(self, default: Quantity | None = None) -> Quantity | None:
+    def min(self, default: QuantityT | None = None) -> QuantityT | None:
         """Return the min value among all phases, or default if they are all `None`.
 
         Args:
@@ -105,16 +105,16 @@ def min(self, default: Quantity | None = None) -> Quantity | None:
         """
         if not any(self):
             return default
-        value: Quantity = functools.reduce(
+        value: QuantityT = functools.reduce(
             lambda x, y: x if x < y else y,
             filter(None, self),
         )
         return value
 
     def map(
         self,
-        function: Callable[[Quantity], Quantity],
-        default: Quantity | None = None,
+        function: Callable[[QuantityT], QuantityT],
+        default: QuantityT | None = None,
     ) -> Self:
         """Apply the given function on each of the phase values and return the result.
 

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

@@ -82,7 +82,7 @@ def __init__(
         self._output_type = output_type
 
     async def _synchronize_metric_timestamps(
-        self, metrics: Set[asyncio.Task[Optional[Sample]]]
+        self, metrics: Set[asyncio.Task[Optional[Sample[Quantity]]]]
     ) -> datetime:
         """Synchronize the metric streams.
 
@@ -129,7 +129,7 @@ async def _synchronize_metric_timestamps(
         self._first_run = False
         return latest_ts
 
-    async def apply(self) -> Sample:
+    async def apply(self) -> Sample[Quantity]:
         """Fetch the latest metrics, apply the formula once and return the result.
 
         Returns:
@@ -304,7 +304,7 @@ def __init__(
         self._name: str = builder.name
         self._builder = builder
         self._output_type = output_type
-        self._channel = Broadcast[Sample](self._name)
+        self._channel = Broadcast[Sample[Quantity]](self._name)
 
     async def _run(self) -> None:
         await self._builder.subscribe()
@@ -328,7 +328,7 @@ async def _run(self) -> None:
 
     def new_receiver(
         self, name: Optional[str] = None, max_size: int = 50
-    ) -> Receiver[Sample]:
+    ) -> Receiver[Sample[Quantity]]:
         """Create a new receiver that streams the output of the formula engine.
 
         Args:
@@ -371,7 +371,7 @@ def __init__(
         self._higher_order_builder = HigherOrderFormulaBuilder3Phase
         self._name: str = name
         self._output_type = output_type
-        self._channel = Broadcast[Sample3Phase](self._name)
+        self._channel = Broadcast[Sample3Phase[Quantity]](self._name)
         self._task: asyncio.Task[None] | None = None
         self._streams: tuple[
             FormulaEngine, FormulaEngine, FormulaEngine
@@ -402,7 +402,7 @@ async def _run(self) -> None:
 
     def new_receiver(
         self, name: Optional[str] = None, max_size: int = 50
-    ) -> Receiver[Sample3Phase]:
+    ) -> Receiver[Sample3Phase[Quantity]]:
         """Create a new receiver that streams the output of the formula engine.
 
         Args:
@@ -493,7 +493,7 @@ def push_oper(self, oper: str) -> None:
     def push_metric(
         self,
         name: str,
-        data_stream: Receiver[Sample],
+        data_stream: Receiver[Sample[Quantity]],
         nones_are_zeros: bool,
     ) -> None:
         """Push a metric receiver into the engine.
@@ -562,7 +562,9 @@ def push_clipper(self, min_value: float | None, max_value: float | None) -> None
         """
         self._steps.append(Clipper(min_value, max_value))
 
-    def push_average(self, metrics: List[Tuple[str, Receiver[Sample], bool]]) -> None:
+    def push_average(
+        self, metrics: List[Tuple[str, Receiver[Sample[Quantity]], bool]]
+    ) -> None:
         """Push an average calculator into the engine.
 
         Args:

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

@@ -11,6 +11,7 @@
 from frequenz.channels import Receiver
 
 from .. import Sample
+from .._quantities import Quantity
 from ._exceptions import FormulaEngineError
 
 
@@ -263,7 +264,7 @@ class MetricFetcher(FormulaStep):
     """A formula step for fetching a value from a metric Receiver."""
 
     def __init__(
-        self, name: str, stream: Receiver[Sample], nones_are_zeros: bool
+        self, name: str, stream: Receiver[Sample[Quantity]], nones_are_zeros: bool
     ) -> None:
         """Create a `MetricFetcher` instance.
 
@@ -274,10 +275,10 @@ def __init__(
         """
         self._name = name
         self._stream = stream
-        self._next_value: Optional[Sample] = None
+        self._next_value: Optional[Sample[Quantity]] = None
         self._nones_are_zeros = nones_are_zeros
 
-    async def fetch_next(self) -> Optional[Sample]:
+    async def fetch_next(self) -> Optional[Sample[Quantity]]:
         """Fetch the next value from the stream.
 
         To be called before each call to `apply`.
@@ -289,7 +290,7 @@ async def fetch_next(self) -> Optional[Sample]:
         return self._next_value
 
     @property
-    def value(self) -> Optional[Sample]:
+    def value(self) -> Optional[Sample[Quantity]]:
         """Get the next value in the stream.
 
         Returns:

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

@@ -54,7 +54,7 @@ def __init__(  # pylint: disable=too-many-arguments
 
     def _get_resampled_receiver(
         self, component_id: int, metric_id: ComponentMetricId
-    ) -> Receiver[Sample]:
+    ) -> Receiver[Sample[Quantity]]:
         """Get a receiver with the resampled data for the given component id.
 
         Args: