Limit Select to work only on Receivers

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

Author: shsms

src/frequenz/channels/select.py

@@ -1,17 +1,19 @@
 # License: MIT
 # Copyright © 2022 Frequenz Energy-as-a-Service GmbH
 
-"""Select the first among multiple AsyncIterators.
+"""Select the first among multiple Receivers.
 
-Expects AsyncIterator class to raise `StopAsyncIteration`
+Expects Receiver class to raise `StopAsyncIteration`
 exception once no more messages are expected or the channel
 is closed in case of `Receiver` class.
 """
 
 import asyncio
 import logging
 from dataclasses import dataclass
-from typing import Any, AsyncIterator, Dict, List, Optional, Set, TypeVar
+from typing import Any, Dict, List, Optional, Set, TypeVar
+
+from frequenz.channels.base_classes import Receiver
 
 logger = logging.Logger(__name__)
 T = TypeVar("T")
@@ -29,17 +31,17 @@ class _Selected:
 
 
 class Select:
-    """Select the next available message from a group of AsyncIterators.
+    """Select the next available message from a group of Receivers.
 
-    If `Select` was created with more `AsyncIterator` than what are read in
+    If `Select` was created with more `Receiver` than what are read in
     the if-chain after each call to [ready()][frequenz.channels.Select.ready],
-    messages coming in the additional async iterators are dropped, and
+    messages coming in the additional receivers are dropped, and
     a warning message is logged.
 
-    [Receiver][frequenz.channels.Receiver]s also function as `AsyncIterator`.
+    [Receiver][frequenz.channels.Receiver]s also function as `Receiver`.
 
     Example:
-        For example, if there are two async iterators that you want to
+        For example, if there are two receivers that you want to
         simultaneously wait on, this can be done with:
 
         ```python
@@ -58,11 +60,11 @@ class Select:
         ```
     """
 
-    def __init__(self, **kwargs: AsyncIterator[Any]) -> None:
+    def __init__(self, **kwargs: Receiver[Any]) -> None:
         """Create a `Select` instance.
 
         Args:
-            **kwargs: sequence of async iterators
+            **kwargs: sequence of receivers
         """
         self._receivers = kwargs
         self._pending: Set[asyncio.Task[Any]] = set()
@@ -84,10 +86,10 @@ def __del__(self) -> None:
             task.cancel()
 
     async def ready(self) -> bool:
-        """Wait until there is a message in any of the async iterators.
+        """Wait until there is a message in any of the receivers.
 
         Returns `True` if there is a message available, and `False` if all
-        async iterators have closed.
+        receivers have closed.
 
         Returns:
             Whether there are further messages or not.
@@ -102,7 +104,7 @@ async def ready(self) -> bool:
                 self._ready_count = 0
                 self._prev_ready_count = 0
                 logger.warning(
-                    "Select.ready() dropped data from async iterator(s): %s, "
+                    "Select.ready() dropped data from receiver(s): %s, "
                     "because no messages have been fetched since the last call to ready().",
                     dropped_names,
                 )
@@ -127,7 +129,7 @@ async def ready(self) -> bool:
                 result = item.result()
             self._ready_count += 1
             self._result[name] = _Selected(result)
-            # if channel or AsyncIterator is closed
+            # if channel or Receiver is closed
             # don't add a task for it again.
             if result is None:
                 continue
@@ -138,13 +140,13 @@ async def ready(self) -> bool:
         return True
 
     def __getattr__(self, name: str) -> Optional[Any]:
-        """Return the latest unread message from a `AsyncIterator`, if available.
+        """Return the latest unread message from a `Receiver`, if available.
 
         Args:
             name: Name of the channel.
 
         Returns:
-            Latest unread message for the specified `AsyncIterator`, or `None`.
+            Latest unread message for the specified `Receiver`, or `None`.
 
         Raises:
             KeyError: when the name was not specified when creating the

tests/test_select.py

@@ -4,67 +4,34 @@
 """Tests for the Select implementation."""
 
 import asyncio
-from asyncio import Queue
-from typing import List, Optional
+from typing import List
 
 from frequenz.channels import Anycast, Select, Sender
 
 
-class AsyncIterable:
-    """Example AsyncIterator class"""
-
-    def __init__(self) -> None:
-        self.queue: "Queue[int]" = Queue()
-        self.done = False
-
-    def __aiter__(self) -> "AsyncIterable":
-        return self
-
-    async def __anext__(self) -> Optional[int]:
-        if not self.queue.empty():
-            return self.queue.get_nowait()
-        if self.done:
-            raise StopAsyncIteration
-        msg = await self.queue.get()
-        return msg
-
-    async def add(self, msg: int) -> bool:
-        """Adds object to iterator"""
-        await self.queue.put(msg)
-
-        return True
-
-
 async def test_select() -> None:
     """Ensure select receives messages in order."""
     chan1 = Anycast[int]()
     chan2 = Anycast[int]()
     chan3 = Anycast[int]()
 
-    async def send(
-        ch1: Sender[int], ch2: Sender[int], ch3: Sender[int], queue: AsyncIterable
-    ) -> None:
+    async def send(ch1: Sender[int], ch2: Sender[int], ch3: Sender[int]) -> None:
         for ctr in range(5):
             await ch1.send(ctr + 1)
             await ch2.send(ctr + 101)
             await ch3.send(ctr + 201)
-            await queue.add(ctr + 301)
         await chan1.close()
         await ch2.send(1000)
         await chan2.close()
         await chan3.close()
-        queue.done = True
-
-    queue = AsyncIterable()
 
     senders = asyncio.create_task(
-        send(chan1.get_sender(), chan2.get_sender(), chan3.get_sender(), queue),
+        send(chan1.get_sender(), chan2.get_sender(), chan3.get_sender()),
     )
     select = Select(
         ch1=chan1.get_receiver(),
         ch2=chan2.get_receiver(),
         ch3=chan3.get_receiver(),
-        ch4=queue,
     )
 
     # only check for messages from all iterators but `ch3`.
@@ -82,32 +49,21 @@ async def send(
                 results.append(val)
             else:
                 results.append(-2)
-        elif item := select.ch4:
-            if val := item.inner:
-                results.append(val)
-            else:
-                results.append(-4)
     await senders
 
     expected_results = [
         1,
         101,
-        301,
         2,
         102,
-        302,
         3,
         103,
-        303,
         4,
         104,
-        304,
         5,
         105,
-        305,
         -1,  # marks end of messages from channel 1
         1000,
-        -4,  # marks end of messages from channel 4
         -2,  # marks end of messages from channel 2
     ]
     assert results == expected_results