Add support/tests for None-value channels (#83)

Author: sahas-subramanian-frequenz

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, Type
 
 from ._base_classes import Receiver as BaseReceiver
 from ._base_classes import Sender as BaseSender
@@ -151,6 +151,10 @@ async def send(self, msg: T) -> None:
             self._chan.recv_cv.notify(1)
 
 
+class _Empty:
+    """A sentinel value to indicate that a value has not been set."""
+
+
 class Receiver(BaseReceiver[T]):
     """A receiver to receive messages from an Anycast channel.
 
@@ -165,7 +169,7 @@ def __init__(self, chan: Anycast[T]) -> None:
             chan: A reference to the channel that this receiver belongs to.
         """
         self._chan = chan
-        self._next: Optional[T] = None
+        self._next: T | Type[_Empty] = _Empty
 
     async def ready(self) -> bool:
         """Wait until the receiver is ready with a value or an error.
@@ -179,7 +183,7 @@ async def ready(self) -> bool:
             Whether the receiver is still active.
         """
         # if a message is already ready, then return immediately.
-        if self._next is not None:
+        if self._next is not _Empty:
             return True
 
         while len(self._chan.deque) == 0:
@@ -202,12 +206,15 @@ def consume(self) -> T:
             ReceiverStoppedError: if the receiver stopped producing messages.
             ReceiverError: if there is some problem with the receiver.
         """
-        if self._next is None and self._chan.closed:
+        if self._next is _Empty and self._chan.closed:
             raise ReceiverStoppedError(self) from ChannelClosedError(self._chan)
 
         assert (
-            self._next is not None
+            self._next is not _Empty
         ), "`consume()` must be preceeded by a call to `ready()`"
-        next_val = self._next
-        self._next = None
+        # mypy doesn't understand that the assert above ensures that self._next is not
+        # _Sentinel.  So we have to use a type ignore here.
+        next_val: T = self._next  # type: ignore[assignment]
+        self._next = _Empty
+
         return next_val

tests/test_anycast.py

@@ -3,6 +3,8 @@
 
 """Tests for the Channel implementation."""
 
+from __future__ import annotations
+
 import asyncio
 
 import pytest
@@ -147,6 +149,23 @@ async def test_anycast_full() -> None:
         assert False
 
 
+async def test_anycast_none_values() -> None:
+    """Ensure None values can be sent and received."""
+    acast: Anycast[int | None] = Anycast()
+
+    sender = acast.new_sender()
+    receiver = acast.new_receiver()
+
+    await sender.send(5)
+    assert await receiver.receive() == 5
+
+    await sender.send(None)
+    assert await receiver.receive() is None
+
+    await sender.send(10)
+    assert await receiver.receive() == 10
+
+
 async def test_anycast_async_iterator() -> None:
     """Check that the anycast receiver works as an async iterator."""
     acast: Anycast[str] = Anycast()

tests/test_broadcast.py

@@ -3,6 +3,8 @@
 
 """Tests for the Broadcast implementation."""
 
+from __future__ import annotations
+
 import asyncio
 from typing import Tuple
 
@@ -69,6 +71,23 @@ async def update_tracker_on_receive(receiver_id: int, recv: Receiver[int]) -> No
     assert actual_sum == expected_sum
 
 
+async def test_broadcast_none_values() -> None:
+    """Ensure None values can be sent and received."""
+    bcast: Broadcast[int | None] = Broadcast("any_channel")
+
+    sender = bcast.new_sender()
+    receiver = bcast.new_receiver()
+
+    await sender.send(5)
+    assert await receiver.receive() == 5
+
+    await sender.send(None)
+    assert await receiver.receive() is None
+
+    await sender.send(10)
+    assert await receiver.receive() == 10
+
+
 async def test_broadcast_after_close() -> None:
     """Ensure closed channels can't get new messages."""
     bcast: Broadcast[int] = Broadcast("meter_5")