Update ready and consume methods to raise ChannelClosedError

... instead of `StopAsyncIteration`

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

Author: shsms

src/frequenz/channels/anycast.py

@@ -9,6 +9,7 @@
 from collections import deque
 from typing import Deque, Generic, Optional
 
+from frequenz.channels.base_classes import ChannelClosedError
 from frequenz.channels.base_classes import Receiver as BaseReceiver
 from frequenz.channels.base_classes import Sender as BaseSender
 from frequenz.channels.base_classes import T
@@ -168,15 +169,15 @@ async def ready(self) -> None:
         """Wait until the receiver is ready with a value.
 
         Raises:
-            StopAsyncIteration: if the underlying channel is closed.
+            ChannelClosedError: if the underlying channel is closed.
         """
         # if a message is already ready, then return immediately.
         if self._next is not None:
             return
 
         while len(self._chan.deque) == 0:
             if self._chan.closed:
-                raise StopAsyncIteration()
+                raise ChannelClosedError()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
         self._next = self._chan.deque.popleft()
@@ -186,9 +187,6 @@ async def ready(self) -> None:
     def consume(self) -> T:
         """Return the latest value once `ready()` is complete.
 
-        Raises:
-            EOFError: When called before a call to `ready()` finishes.
-
         Returns:
             The next value that was received.
         """

src/frequenz/channels/base_classes.py

@@ -69,8 +69,11 @@ async def __anext__(self) -> T:
         Raises:
             StopAsyncIteration: if the underlying channel is closed.
         """
-        await self.ready()
-        return self.consume()
+        try:
+            await self.ready()
+            return self.consume()
+        except ChannelClosedError as exc:
+            raise StopAsyncIteration() from exc
 
     @abstractmethod
     async def ready(self) -> None:
@@ -80,7 +83,7 @@ async def ready(self) -> None:
         `consume()`.
 
         Raises:
-            StopAsyncIteration: if the underlying channel is closed.
+            ChannelClosedError: if the underlying channel is closed.
         """
 
     @abstractmethod
@@ -93,7 +96,7 @@ def consume(self) -> T:
             The next value received.
 
         Raises:
-            StopAsyncIteration: if the underlying channel is closed.
+            ChannelClosedError: if the underlying channel is closed.
         """
 
     def __aiter__(self) -> Receiver[T]:

src/frequenz/channels/broadcast.py

@@ -12,7 +12,7 @@
 from typing import Deque, Dict, Generic, Optional
 from uuid import UUID, uuid4
 
-from frequenz.channels.base_classes import BufferedReceiver
+from frequenz.channels.base_classes import BufferedReceiver, ChannelClosedError
 from frequenz.channels.base_classes import Peekable as BasePeekable
 from frequenz.channels.base_classes import Sender as BaseSender
 from frequenz.channels.base_classes import T
@@ -253,8 +253,8 @@ async def ready(self) -> None:
         """Wait until the receiver is ready with a value.
 
         Raises:
-            EOFError: When called before a call to `ready()` finishes.
-            StopAsyncIteration: if the underlying channel is closed.
+            EOFError: if this receiver is no longer active.
+            ChannelClosedError: if the underlying channel is closed.
         """
         if not self._active:
             raise EOFError("This receiver is no longer active.")
@@ -265,7 +265,7 @@ async def ready(self) -> None:
         # consumed, then we return immediately.
         while len(self._q) == 0:
             if self._chan.closed:
-                raise StopAsyncIteration()
+                raise ChannelClosedError()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
 

src/frequenz/channels/merge.py

@@ -7,7 +7,7 @@
 from collections import deque
 from typing import Any, Deque, Set
 
-from frequenz.channels.base_classes import Receiver, T
+from frequenz.channels.base_classes import ChannelClosedError, Receiver, T
 
 
 class Merge(Receiver[T]):
@@ -48,7 +48,7 @@ async def ready(self) -> None:
         """Wait until the receiver is ready with a value.
 
         Raises:
-            StopAsyncIteration: if the underlying channel is closed.
+            ChannelClosedError: if the underlying channel is closed.
         """
         # we use a while loop to continue to wait for new data, in case the
         # previous `wait` completed because a channel was closed.
@@ -58,7 +58,7 @@ async def ready(self) -> None:
                 return
 
             if len(self._pending) == 0:
-                raise StopAsyncIteration()
+                raise ChannelClosedError()
             done, self._pending = await asyncio.wait(
                 self._pending, return_when=asyncio.FIRST_COMPLETED
             )

src/frequenz/channels/merge_named.py

@@ -7,7 +7,7 @@
 from collections import deque
 from typing import Any, Deque, Set, Tuple
 
-from frequenz.channels.base_classes import Receiver, T
+from frequenz.channels.base_classes import ChannelClosedError, Receiver, T
 
 
 class MergeNamed(Receiver[Tuple[str, T]]):
@@ -35,7 +35,7 @@ async def ready(self) -> None:
         """Wait until there's a message in any of the channels.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
+            ChannelClosedError: when all the channels are closed.
         """
         # we use a while loop to continue to wait for new data, in case the
         # previous `wait` completed because a channel was closed.
@@ -45,7 +45,7 @@ async def ready(self) -> None:
                 return
 
             if len(self._pending) == 0:
-                raise StopAsyncIteration()
+                raise ChannelClosedError()
             done, self._pending = await asyncio.wait(
                 self._pending, return_when=asyncio.FIRST_COMPLETED
             )

src/frequenz/channels/select.py

@@ -13,7 +13,7 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Set, TypeVar
 
-from frequenz.channels.base_classes import Receiver
+from frequenz.channels.base_classes import ChannelClosedError, Receiver
 
 logger = logging.Logger(__name__)
 T = TypeVar("T")
@@ -147,7 +147,7 @@ async def ready(self) -> bool:
         for item in done:
             name = item.get_name()
             recv = self._receivers[name]
-            if isinstance(item.exception(), StopAsyncIteration):
+            if isinstance(item.exception(), ChannelClosedError):
                 result = None
             else:
                 result = recv
@@ -157,8 +157,7 @@ async def ready(self) -> bool:
             # don't add a task for it again.
             if result is None:
                 continue
-            ready = recv.ready()
-            self._pending.add(asyncio.create_task(ready, name=name))
+            self._pending.add(asyncio.create_task(recv.ready(), name=name))
         return True
 
     def __getattr__(self, name: str) -> Optional[Any]:

src/frequenz/channels/utils/file_watcher.py

@@ -10,7 +10,7 @@
 from watchfiles import Change, awatch
 from watchfiles.main import FileChange
 
-from frequenz.channels.base_classes import Receiver
+from frequenz.channels.base_classes import ChannelClosedError, Receiver
 
 
 class EventType(Enum):
@@ -83,7 +83,7 @@ def consume(self) -> pathlib.Path:
         """Return the latest change once `ready` is complete.
 
         Raises:
-            StopAsyncIteration: When the channel is closed.
+            ChannelClosedError: When the channel is closed.
 
         Returns:
             The next change that was received.
@@ -92,7 +92,7 @@ def consume(self) -> pathlib.Path:
         change = self._changes.pop()
         # Tuple of (Change, path) returned by watchfiles
         if change is None or len(change) != 2:
-            raise StopAsyncIteration()
+            raise ChannelClosedError()
         _, path_str = change
         path = pathlib.Path(path_str)
         return path

src/frequenz/channels/utils/timer.py

@@ -7,7 +7,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import Optional
 
-from frequenz.channels.base_classes import Receiver
+from frequenz.channels.base_classes import ChannelClosedError, Receiver
 
 
 class Timer(Receiver[datetime]):
@@ -81,7 +81,7 @@ async def ready(self) -> None:
         """Return the current time (in UTC) once the next tick is due.
 
         Raises:
-            StopAsyncIteration: if [stop()][frequenz.channels.Timer.stop] has been
+            ChannelClosedError: if [stop()][frequenz.channels.Timer.stop] has been
                 called on the timer.
 
         Returns:
@@ -94,7 +94,7 @@ async def ready(self) -> None:
             return
 
         if self._stopped:
-            raise StopAsyncIteration()
+            raise ChannelClosedError()
         now = datetime.now(timezone.utc)
         diff = self._next_msg_time - now
         while diff.total_seconds() > 0: