Adapt code to use the new ChannelRegistry

Note that we needed to include a hack for the `BatteryPool` because the
hierarchy of `_Report`/`Report`/`BatteryPoolReport` is messed up and the
types can only be forced to be compatible with a cast.

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

benchmarks/timeseries/benchmark_datasourcing.py

@@ -108,7 +108,9 @@ async def consume(channel: Receiver[Any]) -> None:
                 "current_phase_requests", evc_id, component_metric_id, None
             )
 
-            recv_channel = channel_registry.new_receiver(request.get_channel_name())
+            recv_channel = channel_registry.get_or_create(
+                ComponentMetricRequest, request.get_channel_name()
+            ).new_receiver()
 
             await request_sender.send(request)
             consume_tasks.append(asyncio.create_task(consume(recv_channel)))

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

@@ -325,7 +325,9 @@ def _get_metric_senders(
             (
                 self._get_data_extraction_method(category, metric),
                 [
-                    self._registry.new_sender(request.get_channel_name())
+                    self._registry.get_or_create(
+                        Sample[Quantity], request.get_channel_name()
+                    ).new_sender()
                     for request in req_list
                 ],
             )
@@ -379,8 +381,7 @@ def process_msg(data: Any) -> None:
 
         await asyncio.gather(
             *[
-                # pylint: disable=protected-access
-                self._registry._close_channel(r.get_channel_name())
+                self._registry.close_and_remove(r.get_channel_name())
                 for requests in self._req_streaming_metrics[comp_id].values()
                 for r in requests
             ]

src/frequenz/sdk/actor/_power_managing/_power_managing_actor.py

@@ -211,14 +211,16 @@ async def _run(self) -> None:
 
                 if component_ids not in self._subscriptions:
                     self._subscriptions[component_ids] = {
-                        priority: self._channel_registry.new_sender(
-                            sub.get_channel_name()
-                        )
+                        priority: self._channel_registry.get_or_create(
+                            _Report, sub.get_channel_name()
+                        ).new_sender()
                     }
                 elif priority not in self._subscriptions[component_ids]:
                     self._subscriptions[component_ids][
                         priority
-                    ] = self._channel_registry.new_sender(sub.get_channel_name())
+                    ] = self._channel_registry.get_or_create(
+                        _Report, sub.get_channel_name()
+                    ).new_sender()
 
                 if sub.component_ids not in self._bound_tracker_tasks:
                     self._add_bounds_tracker(sub.component_ids)

src/frequenz/sdk/actor/_resampling.py

@@ -77,11 +77,15 @@ async def _subscribe(self, request: ComponentMetricRequest) -> None:
         )
         data_source_channel_name = data_source_request.get_channel_name()
         await self._data_sourcing_request_sender.send(data_source_request)
-        receiver = self._channel_registry.new_receiver(data_source_channel_name)
+        receiver = self._channel_registry.get_or_create(
+            Sample[Quantity], data_source_channel_name
+        ).new_receiver()
 
         # This is a temporary hack until the Sender implementation uses
         # exceptions to report errors.
-        sender = self._channel_registry.new_sender(request_channel_name)
+        sender = self._channel_registry.get_or_create(
+            Sample[Quantity], request_channel_name
+        ).new_sender()
 
         async def sink_adapter(sample: Sample[Quantity]) -> None:
             await sender.send(sample)

src/frequenz/sdk/timeseries/_grid_frequency.py

@@ -15,7 +15,7 @@
 from ..microgrid import connection_manager
 from ..microgrid.component import Component, ComponentCategory, ComponentMetricId
 from ..timeseries._base_types import Sample
-from ..timeseries._quantities import Frequency
+from ..timeseries._quantities import Frequency, Quantity
 
 if TYPE_CHECKING:
     # Imported here to avoid a circular import.
@@ -95,9 +95,9 @@ def new_receiver(self) -> Receiver[Sample[Frequency]]:
         Returns:
             A receiver that will receive grid frequency samples.
         """
-        receiver = self._channel_registry.new_receiver(
-            self._component_metric_request.get_channel_name()
-        )
+        receiver = self._channel_registry.get_or_create(
+            Sample[Quantity], self._component_metric_request.get_channel_name()
+        ).new_receiver()
 
         if not self._task:
             self._task = asyncio.create_task(self._send_request())

src/frequenz/sdk/timeseries/_voltage_streamer.py

@@ -118,7 +118,9 @@ def new_receiver(self) -> Receiver[Sample3Phase[Voltage]]:
             A receiver that will receive the phase-to-neutral voltage as a
             3-phase sample.
         """
-        receiver = self._channel_registry.new_receiver(self._channel_key)
+        receiver = self._channel_registry.get_or_create(
+            Sample3Phase[Voltage], self._channel_key
+        ).new_receiver()
 
         if not self._task:
             self._task = asyncio.create_task(self._send_request())
@@ -148,10 +150,14 @@ async def _send_request(self) -> None:
             await self._resampler_subscription_sender.send(req)
 
             phases_rx.append(
-                self._channel_registry.new_receiver(req.get_channel_name())
+                self._channel_registry.get_or_create(
+                    Sample[Quantity], req.get_channel_name()
+                ).new_receiver()
             )
 
-        sender = self._channel_registry.new_sender(self._channel_key)
+        sender = self._channel_registry.get_or_create(
+            Sample3Phase[Voltage], self._channel_key
+        ).new_sender()
 
         _logger.debug(
             "Sent request for fetching voltage from: %s", self._source_component

src/frequenz/sdk/timeseries/battery_pool/_battery_pool.py

@@ -12,6 +12,7 @@
 import uuid
 from collections import abc
 from datetime import timedelta
+from typing import cast
 
 from ... import timeseries
 from ..._internal._channels import ReceiverFetcher
@@ -386,12 +387,14 @@ def power_status(self) -> ReceiverFetcher[BatteryPoolReport]:
         ] = asyncio.create_task(
             self._battery_pool._power_manager_bounds_subscription_sender.send(sub)
         )
-        self._battery_pool._channel_registry.set_resend_latest(
-            sub.get_channel_name(), True
-        )
-        return self._battery_pool._channel_registry.new_receiver_fetcher(
-            sub.get_channel_name()
+        channel = self._battery_pool._channel_registry.get_or_create(
+            _power_managing._Report, sub.get_channel_name()
         )
+        channel.resend_latest = True
+
+        # More details on why the cast is needed here:
+        # https://github.com/frequenz-floss/frequenz-sdk-python/issues/823
+        return cast(ReceiverFetcher[BatteryPoolReport], channel)
 
     @property
     def _system_power_bounds(self) -> ReceiverFetcher[SystemBounds]:

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

@@ -279,7 +279,9 @@ async def resampler_subscribe(
                 start_time=None,
             )
             await self._resampler_subscription_sender.send(request)
-            return self._channel_registry.new_receiver(request.get_channel_name())
+            return self._channel_registry.get_or_create(
+                Sample[Quantity], request.get_channel_name()
+            ).new_receiver()
 
         return (
             await resampler_subscribe(ComponentMetricId.CURRENT_PHASE_1),

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

@@ -12,7 +12,7 @@
 
 from ...microgrid.component import ComponentMetricId
 from .. import Sample
-from .._quantities import QuantityT
+from .._quantities import Quantity, QuantityT
 from ._formula_engine import FormulaBuilder, FormulaEngine
 from ._tokenizer import Tokenizer, TokenType
 
@@ -74,7 +74,9 @@ def _get_resampled_receiver(
 
         request = ComponentMetricRequest(self._namespace, component_id, metric_id, None)
         self._resampler_requests.append(request)
-        return self._channel_registry.new_receiver(request.get_channel_name())
+        return self._channel_registry.get_or_create(
+            Sample[QuantityT], request.get_channel_name()
+        ).new_receiver()
 
     async def subscribe(self) -> None:
         """Subscribe to all resampled component metric streams."""

tests/actor/test_data_sourcing.py

@@ -13,6 +13,7 @@
 )
 from frequenz.sdk.microgrid import connection_manager
 from frequenz.sdk.microgrid.component import ComponentMetricId
+from frequenz.sdk.timeseries import Quantity, Sample
 from tests.microgrid import mock_api
 
 # pylint: disable=no-member
@@ -59,21 +60,25 @@ async def test_data_sourcing_actor(self) -> None:
             active_power_request = ComponentMetricRequest(
                 "test-namespace", 4, ComponentMetricId.ACTIVE_POWER, None
             )
-            active_power_recv = registry.new_receiver(
-                active_power_request.get_channel_name()
-            )
+            active_power_recv = registry.get_or_create(
+                Sample[Quantity], active_power_request.get_channel_name()
+            ).new_receiver()
             await req_sender.send(active_power_request)
 
             soc_request = ComponentMetricRequest(
                 "test-namespace", 9, ComponentMetricId.SOC, None
             )
-            soc_recv = registry.new_receiver(soc_request.get_channel_name())
+            soc_recv = registry.get_or_create(
+                Sample[Quantity], soc_request.get_channel_name()
+            ).new_receiver()
             await req_sender.send(soc_request)
 
             soc2_request = ComponentMetricRequest(
                 "test-namespace", 9, ComponentMetricId.SOC, None
             )
-            soc2_recv = registry.new_receiver(soc2_request.get_channel_name())
+            soc2_recv = registry.get_or_create(
+                Sample[Quantity], soc2_request.get_channel_name()
+            ).new_receiver()
             await req_sender.send(soc2_request)
 
             for _ in range(3):