Reverse call order of __anext__ and receive methods in Receiver

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

Author: shsms

src/frequenz/channels/anycast.py

@@ -7,7 +7,7 @@
 
 from asyncio import Condition
 from collections import deque
-from typing import Deque, Generic, Optional
+from typing import Deque, Generic
 
 from frequenz.channels.base_classes import Receiver as BaseReceiver
 from frequenz.channels.base_classes import Sender as BaseSender
@@ -163,19 +163,22 @@ def __init__(self, chan: Anycast[T]) -> None:
         """
         self._chan = chan
 
-    async def receive(self) -> Optional[T]:
+    async def __anext__(self) -> T:
         """Receive a message from the channel.
 
         Waits for an message to become available, and returns that message.
         When there are multiple receivers for the channel, only one receiver
         will receive each message.
 
+        Raises:
+            StopAsyncIteration: When the channel is closed.
+
         Returns:
-            `None`, if the channel is closed, a message otherwise.
+            The received message.
         """
         while len(self._chan.deque) == 0:
             if self._chan.closed:
-                return None
+                raise StopAsyncIteration()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
         ret = self._chan.deque.popleft()

src/frequenz/channels/base_classes.py

@@ -32,11 +32,14 @@ class Receiver(ABC, Generic[T]):
     """A channel Receiver."""
 
     @abstractmethod
-    async def receive(self) -> Optional[T]:
-        """Receive a message from the channel.
+    async def __anext__(self) -> T:
+        """Await the next value in the async iteration over received values.
 
         Returns:
-            `None`, if the channel is closed, a message otherwise.
+            The next value received.
+
+        Raises:
+            StopAsyncIteration: if the underlying channel is closed.
         """
 
     def __aiter__(self) -> Receiver[T]:
@@ -47,19 +50,16 @@ def __aiter__(self) -> Receiver[T]:
         """
         return self
 
-    async def __anext__(self) -> T:
-        """Await the next value in the async iteration over received values.
+    async def receive(self) -> Optional[T]:
+        """Receive a message from the channel.
 
         Returns:
-            The next value received.
-
-        Raises:
-            StopAsyncIteration: if we receive `None`, i.e. if the underlying
-                channel is closed.
+            The received message.
         """
-        received = await self.receive()
-        if received is None:
-            raise StopAsyncIteration
+        try:
+            received = await self.__anext__()  # pylint: disable=unnecessary-dunder-call
+        except StopAsyncIteration:
+            return None
         return received
 
     def map(self, call: Callable[[T], U]) -> Receiver[U]:
@@ -136,13 +136,11 @@ def __init__(self, recv: Receiver[T], transform: Callable[[T], U]) -> None:
         self._recv = recv
         self._transform = transform
 
-    async def receive(self) -> Optional[U]:
+    async def __anext__(self) -> U:
         """Return a transformed message received from the input channel.
 
         Returns:
-            `None`, if the channel is closed, a message otherwise.
+            A transformed message.
         """
-        msg = await self._recv.receive()
-        if msg is None:
-            return None
+        msg = await self._recv.__anext__()
         return self._transform(msg)

src/frequenz/channels/bidirectional.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from typing import Generic, Optional
+from typing import Generic
 
 from frequenz.channels.base_classes import Receiver, Sender, T, U
 from frequenz.channels.broadcast import Broadcast
@@ -82,10 +82,10 @@ async def send(self, msg: T) -> bool:
         """
         return await self._sender.send(msg)
 
-    async def receive(self) -> Optional[U]:
+    async def __anext__(self) -> U:
         """Receive a value from the other side.
 
         Returns:
             Received value, or `None` if the channels are closed.
         """
-        return await self._receiver.receive()
+        return await self._receiver.__anext__()

src/frequenz/channels/broadcast.py

@@ -249,7 +249,7 @@ def __len__(self) -> int:
         """
         return len(self._q)
 
-    async def receive(self) -> Optional[T]:
+    async def __anext__(self) -> T:
         """Receive a message from the Broadcast channel.
 
         Waits until there are messages available in the channel and returns
@@ -261,6 +261,7 @@ async def receive(self) -> Optional[T]:
         `EOFError`.
 
         Raises:
+            StopAsyncIteration: When the channel is closed.
             EOFError: when the receiver has been converted into a `Peekable`.
 
         Returns:
@@ -271,7 +272,7 @@ async def receive(self) -> Optional[T]:
 
         while len(self._q) == 0:
             if self._chan.closed:
-                return None
+                raise StopAsyncIteration()
             async with self._chan.recv_cv:
                 await self._chan.recv_cv.wait()
         ret = self._q.popleft()

src/frequenz/channels/merge.py

@@ -5,7 +5,7 @@
 
 import asyncio
 from collections import deque
-from typing import Any, Deque, Optional, Set
+from typing import Any, Deque, Set
 
 from frequenz.channels.base_classes import Receiver, T
 
@@ -34,7 +34,7 @@ def __init__(self, *args: Receiver[T]) -> None:
         """
         self._receivers = {str(id): recv for id, recv in enumerate(args)}
         self._pending: Set[asyncio.Task[Any]] = {
-            asyncio.create_task(recv.receive(), name=name)
+            asyncio.create_task(recv.__anext__(), name=name)
             for name, recv in self._receivers.items()
         }
         self._results: Deque[T] = deque(maxlen=len(self._receivers))
@@ -44,9 +44,12 @@ def __del__(self) -> None:
         for task in self._pending:
             task.cancel()
 
-    async def receive(self) -> Optional[T]:
+    async def __anext__(self) -> T:
         """Wait until there's a message in any of the channels.
 
+        Raises:
+            StopAsyncIteration: When the channel is closed.
+
         Returns:
             The next message that was received, or `None`, if all channels have
                 closed.
@@ -58,17 +61,17 @@ async def receive(self) -> Optional[T]:
                 return self._results.popleft()
 
             if len(self._pending) == 0:
-                return None
+                raise StopAsyncIteration()
             done, self._pending = await asyncio.wait(
                 self._pending, return_when=asyncio.FIRST_COMPLETED
             )
             for item in done:
                 name = item.get_name()
-                result = item.result()
                 # if channel is closed, don't add a task for it again.
-                if result is None:
+                if isinstance(item.exception(), StopAsyncIteration):
                     continue
+                result = item.result()
                 self._results.append(result)
                 self._pending.add(
-                    asyncio.create_task(self._receivers[name].receive(), name=name)
+                    asyncio.create_task(self._receivers[name].__anext__(), name=name)
                 )

src/frequenz/channels/merge_named.py

@@ -5,7 +5,7 @@
 
 import asyncio
 from collections import deque
-from typing import Any, Deque, Optional, Set, Tuple
+from typing import Any, Deque, Set, Tuple
 
 from frequenz.channels.base_classes import Receiver, T
 
@@ -21,7 +21,7 @@ def __init__(self, **kwargs: Receiver[T]) -> None:
         """
         self._receivers = kwargs
         self._pending: Set[asyncio.Task[Any]] = {
-            asyncio.create_task(recv.receive(), name=name)
+            asyncio.create_task(recv.__anext__(), name=name)
             for name, recv in self._receivers.items()
         }
         self._results: Deque[Tuple[str, T]] = deque(maxlen=len(self._receivers))
@@ -31,9 +31,12 @@ def __del__(self) -> None:
         for task in self._pending:
             task.cancel()
 
-    async def receive(self) -> Optional[Tuple[str, T]]:
+    async def __anext__(self) -> Tuple[str, T]:
         """Wait until there's a message in any of the channels.
 
+        Raises:
+            StopAsyncIteration: When the channel is closed.
+
         Returns:
             The next message that was received, or `None`, if all channels have
                 closed.
@@ -45,17 +48,17 @@ async def receive(self) -> Optional[Tuple[str, T]]:
                 return self._results.popleft()
 
             if len(self._pending) == 0:
-                return None
+                raise StopAsyncIteration()
             done, self._pending = await asyncio.wait(
                 self._pending, return_when=asyncio.FIRST_COMPLETED
             )
             for item in done:
                 name = item.get_name()
-                result = item.result()
                 # if channel is closed, don't add a task for it again.
-                if result is None:
+                if isinstance(item.exception(), StopAsyncIteration):
                     continue
+                result = item.result()
                 self._results.append((name, result))
                 self._pending.add(
-                    asyncio.create_task(self._receivers[name].receive(), name=name)
+                    asyncio.create_task(self._receivers[name].__anext__(), name=name)
                 )

src/frequenz/channels/utils/file_watcher.py

@@ -62,9 +62,12 @@ def __del__(self) -> None:
         """
         self._stop_event.set()
 
-    async def receive(self) -> Optional[pathlib.Path]:
+    async def __anext__(self) -> pathlib.Path:
         """Wait for the next file event and return its path.
 
+        Raises:
+            StopAsyncIteration: When the channel is closed.
+
         Returns:
             Path of next file.
         """
@@ -73,7 +76,7 @@ async def receive(self) -> Optional[pathlib.Path]:
             for change in changes:
                 # Tuple of (Change, path) returned by watchfiles
                 if change is None or len(change) != 2:
-                    return None
+                    raise StopAsyncIteration()
 
                 _, path_str = change
                 path = pathlib.Path(path_str)

src/frequenz/channels/utils/timer.py

@@ -5,7 +5,6 @@
 
 import asyncio
 from datetime import datetime, timedelta, timezone
-from typing import Optional
 
 from frequenz.channels.base_classes import Receiver
 
@@ -76,9 +75,12 @@ def stop(self) -> None:
         """
         self._stopped = True
 
-    async def receive(self) -> Optional[datetime]:
+    async def __anext__(self) -> datetime:
         """Return the current time (in UTC) once the next tick is due.
 
+        Raises:
+            StopAsyncIteration: When the channel is closed.
+
         Returns:
             The time of the next tick in UTC or `None` if
                 [stop()][frequenz.channels.Timer.stop] has been called on the
@@ -89,7 +91,7 @@ async def receive(self) -> Optional[datetime]:
               instead of a native datetime object.
         """
         if self._stopped:
-            return None
+            raise StopAsyncIteration()
         now = datetime.now(timezone.utc)
         diff = self._next_msg_time - now
         while diff.total_seconds() > 0: