Implement BatchUnrollReceiver

This turns a `Receiver[ComponentDataBatch]` into a
`Receiver[MetricSample]` by unrolling `ComponentDataBatch` objects
into multiple `MetricSample`s.

Then this receiver is used in the client to return `Receiver`s instead
of iterators for component data streams.

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

Author: shsms

src/frequenz/client/reporting/_batch_unroll_receiver.py

@@ -0,0 +1,72 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""A receiver that unrolls batches of data into individual samples."""
+
+from collections.abc import Iterator
+
+from frequenz.channels import Receiver, ReceiverStoppedError
+from typing_extensions import override
+
+from ._types import ComponentsDataBatch, MetricSample
+
+
+class BatchUnrollReceiver(Receiver[MetricSample]):
+    """Receiver to unroll `ComponentsDataBatch`s into `MetricSample`s."""
+
+    def __init__(self, stream: Receiver[ComponentsDataBatch]) -> None:
+        """Initialize the receiver.
+
+        Args:
+            stream: The stream to receive batches from.
+        """
+        self._stream: Receiver[ComponentsDataBatch] = stream
+        self._batch_iter: Iterator[MetricSample] | None = None
+        self._latest_sample: MetricSample | None = None
+        self._no_more_data: bool = False
+
+    @override
+    async def ready(self) -> bool:
+        """Wait until the next `MetricSample` is ready."""
+        # If ready is called multiple times, we should return the same result
+        # so we don't loose any data
+        if self._latest_sample is not None:
+            return True
+
+        while True:
+            # If we have a batch iterator, try to get the next sample
+            if self._batch_iter is not None:
+                try:
+                    self._latest_sample = next(self._batch_iter)
+                    return True
+                # If the batch is done, set the batch iterator to None
+                except StopIteration:
+                    self._batch_iter = None
+
+            # If we don't have a batch iterator, try to get the next batch
+            try:
+                batch = await anext(self._stream)
+                self._batch_iter = iter(batch)
+            # If the stream is done, return False
+            except StopAsyncIteration:
+                self._no_more_data = True
+                return False
+
+    @override
+    def consume(self) -> MetricSample:
+        """Consume the next `MetricSample`.
+
+        Returns:
+            The next `MetricSample`.
+
+        Raises:
+            ReceiverStoppedError: If the receiver is stopped.
+            RuntimeError: If the receiver is not ready.
+        """
+        sample = self._latest_sample
+        if sample is None:
+            if self._no_more_data:
+                raise ReceiverStoppedError(self)
+            raise RuntimeError("consume called before ready")
+        self._latest_sample = None
+        return sample

src/frequenz/client/reporting/_client.py

@@ -3,11 +3,9 @@
 
 """Client for requests to the Reporting API."""
 
-from collections import abc, namedtuple
-from collections.abc import AsyncIterator, Iterator
-from dataclasses import dataclass
-from datetime import datetime, timedelta, timezone
-from typing import Any, AsyncIterable, cast
+from collections.abc import AsyncIterable, AsyncIterator
+from datetime import datetime, timedelta
+from typing import cast
 
 # pylint: disable=no-name-in-module
 from frequenz.api.common.v1.microgrid.microgrid_pb2 import (
@@ -46,13 +44,15 @@
 )
 from frequenz.api.reporting.v1.reporting_pb2 import TimeFilter as PBTimeFilter
 from frequenz.api.reporting.v1.reporting_pb2_grpc import ReportingStub
+from frequenz.channels import Receiver
 from frequenz.client.base.channel import ChannelOptions
 from frequenz.client.base.client import BaseApiClient
 from frequenz.client.base.exception import ClientNotConnected
 from frequenz.client.base.streaming import GrpcStreamBroadcaster
 from frequenz.client.common.metric import Metric
 from google.protobuf.timestamp_pb2 import Timestamp as PBTimestamp
 
+from ._batch_unroll_receiver import BatchUnrollReceiver
 from ._types import (
     AggregatedMetric,
     ComponentsDataBatch,
@@ -113,7 +113,7 @@ def stub(self) -> ReportingStub:
         return self._stub
 
     # pylint: disable=too-many-arguments
-    async def receive_single_component_data(
+    def receive_single_component_data(
         self,
         *,
         microgrid_id: int,
@@ -124,7 +124,7 @@ async def receive_single_component_data(
         resampling_period: timedelta | None,
         include_states: bool = False,
         include_bounds: bool = False,
-    ) -> AsyncIterator[MetricSample]:
+    ) -> Receiver[MetricSample]:
         """Iterate over the data for a single metric.
 
         Args:
@@ -137,12 +137,10 @@ async def receive_single_component_data(
             include_states: Whether to include the state data.
             include_bounds: Whether to include the bound data.
 
-        Yields:
-            A named tuple with the following fields:
-            * timestamp: The timestamp of the metric sample.
-            * value: The metric value.
+        Returns:
+            A receiver of `MetricSample`s.
         """
-        broadcaster = await self._receive_microgrid_components_data_batch(
+        receiver = self._receive_microgrid_components_data_batch(
             microgrid_components=[(microgrid_id, [component_id])],
             metrics=[metrics] if isinstance(metrics, Metric) else metrics,
             start_time=start_time,
@@ -152,14 +150,10 @@ async def receive_single_component_data(
             include_bounds=include_bounds,
         )
 
-        receiver = broadcaster.new_receiver()
-
-        async for batch in receiver:
-            for entry in batch:
-                yield entry
+        return BatchUnrollReceiver(receiver)
 
     # pylint: disable=too-many-arguments
-    async def receive_microgrid_components_data(
+    def receive_microgrid_components_data(
         self,
         *,
         microgrid_components: list[tuple[int, list[int]]],
@@ -169,7 +163,7 @@ async def receive_microgrid_components_data(
         resampling_period: timedelta | None,
         include_states: bool = False,
         include_bounds: bool = False,
-    ) -> AsyncIterator[MetricSample]:
+    ) -> Receiver[MetricSample]:
         """Iterate over the data for multiple microgrids and components.
 
         Args:
@@ -182,15 +176,10 @@ async def receive_microgrid_components_data(
             include_states: Whether to include the state data.
             include_bounds: Whether to include the bound data.
 
-        Yields:
-            A named tuple with the following fields:
-            * microgrid_id: The microgrid ID.
-            * component_id: The component ID.
-            * metric: The metric name.
-            * timestamp: The timestamp of the metric sample.
-            * value: The metric value.
+        Returns:
+            A receiver of `MetricSample`s.
         """
-        broadcaster = await self._receive_microgrid_components_data_batch(
+        receiver = self._receive_microgrid_components_data_batch(
             microgrid_components=microgrid_components,
             metrics=[metrics] if isinstance(metrics, Metric) else metrics,
             start_time=start_time,
@@ -200,15 +189,11 @@ async def receive_microgrid_components_data(
             include_bounds=include_bounds,
         )
 
-        receiver = broadcaster.new_receiver()
-
-        async for batch in receiver:
-            for entry in batch:
-                yield entry
+        return BatchUnrollReceiver(receiver)
 
     # pylint: disable=too-many-arguments
     # pylint: disable=too-many-locals
-    async def _receive_microgrid_components_data_batch(
+    def _receive_microgrid_components_data_batch(
         self,
         *,
         microgrid_components: list[tuple[int, list[int]]],
@@ -218,9 +203,7 @@ async def _receive_microgrid_components_data_batch(
         resampling_period: timedelta | None,
         include_states: bool = False,
         include_bounds: bool = False,
-    ) -> GrpcStreamBroadcaster[
-        PBReceiveMicrogridComponentsDataStreamResponse, ComponentsDataBatch
-    ]:
+    ) -> Receiver[ComponentsDataBatch]:
         """Return a GrpcStreamBroadcaster for microgrid component data."""
         stream_key = (
             tuple((mid, tuple(cids)) for mid, cids in microgrid_components),
@@ -315,7 +298,7 @@ async def stream_method() -> (
                 retry_strategy=None,
             )
 
-        return self._components_data_streams[stream_key]
+        return self._components_data_streams[stream_key].new_receiver()
 
     # pylint: disable=too-many-arguments
     async def receive_single_sensor_data(