Add take_while and drop_while to Receiver

The `take_while()` method is just an alias for `filter()`, with the
intention to eliminate the ambiguity and make it more readable.

On the other hand, `drop_while()` is the negation of `take_while()` and
provided as a convenience and also to improve readability.

These names are widely popular and used in other programming languages
as well as the Python `itertools` module.

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

Author: llucax

RELEASE_NOTES.md

@@ -2,6 +2,12 @@
 
 ## New Features
 
+- `Receiver`
+
+   * Add `take_while()` as a less ambiguous and more readable alternative to `filter()`.
+   * Add `drop_while()` as a convenience and more readable alternative to `filter()` with a negated predicate.
+   * The usage of `filter()` is discouraged in favor of `take_while()` and `drop_while()`.
+
 ### Experimental
 
 - A new predicate, `OnlyIfPrevious`, to `filter()` messages based on the previous message.

src/frequenz/channels/_receiver.py

@@ -58,21 +58,23 @@
 # Message Filtering
 
 If you need to filter the received messages, receivers provide a
-[`filter()`][frequenz.channels.Receiver.filter] method to easily do so:
+[`take_while()`][frequenz.channels.Receiver.take_while] and a
+[`drop_while()`][frequenz.channels.Receiver.drop_while]
+method to easily do so:
 
 ```python show_lines="6:"
 from frequenz.channels import Anycast
 
 channel = Anycast[int](name="test-channel")
 receiver = channel.new_receiver()
 
-async for message in receiver.filter(lambda x: x % 2 == 0):
+async for message in receiver.take_while(lambda x: x % 2 == 0):
     print(message)  # Only even numbers will be printed
 ```
 
 As with [`map()`][frequenz.channels.Receiver.map],
-[`filter()`][frequenz.channels.Receiver.filter] returns a new full receiver, so you can
-use it in any of the ways described above.
+[`take_while()`][frequenz.channels.Receiver.take_while] returns a new full receiver, so
+you can use it in any of the ways described above.
 
 # Error Handling
 
@@ -280,6 +282,11 @@ def filter(
     ) -> Receiver[FilteredMessageT_co]:
         """Apply a type guard on the messages on a receiver.
 
+        Tip:
+            It is recommended to use the
+            [`take_while()`][frequenz.channels.Receiver.take_while] method instead of
+            this one, as it makes the intention more clear.
+
         Tip:
             The returned receiver type won't have all the methods of the original
             receiver. If you need to access methods of the original receiver that are
@@ -301,6 +308,11 @@ def filter(
     ) -> Receiver[ReceiverMessageT_co]:
         """Apply a filter function on the messages on a receiver.
 
+        Tip:
+            It is recommended to use the
+            [`take_while()`][frequenz.channels.Receiver.take_while] method instead of
+            this one, as it makes the intention more clear.
+
         Tip:
             The returned receiver type won't have all the methods of the original
             receiver. If you need to access methods of the original receiver that are
@@ -326,6 +338,11 @@ def filter(
     ) -> Receiver[ReceiverMessageT_co] | Receiver[FilteredMessageT_co]:
         """Apply a filter function on the messages on a receiver.
 
+        Tip:
+            It is recommended to use the
+            [`take_while()`][frequenz.channels.Receiver.take_while] method instead of
+            this one, as it makes the intention more clear.
+
         Note:
             You can pass a [type guard][typing.TypeGuard] as the filter function to
             narrow the type of the messages that pass the filter.
@@ -345,6 +362,117 @@ def filter(
         """
         return _Filter(receiver=self, filter_function=filter_function)
 
+    @overload
+    def take_while(
+        self,
+        predicate: Callable[[ReceiverMessageT_co], TypeGuard[FilteredMessageT_co]],
+        /,
+    ) -> Receiver[FilteredMessageT_co]:
+        """Take only the messages that fulfill a predicate, narrowing the type.
+
+        The returned receiver will only receive messages that fulfill the predicate
+        (evaluates to `True`), and will drop messages that don't.
+
+        Tip:
+            The returned receiver type won't have all the methods of the original
+            receiver. If you need to access methods of the original receiver that are
+            not part of the `Receiver` interface you should save a reference to the
+            original receiver and use that instead.
+
+        Args:
+            predicate: The predicate to be applied on incoming messages to
+                determine if they should be taken.
+
+        Returns:
+            A new receiver that only receives messages that fulfill the predicate.
+        """
+        ...  # pylint: disable=unnecessary-ellipsis
+
+    @overload
+    def take_while(
+        self, predicate: Callable[[ReceiverMessageT_co], bool], /
+    ) -> Receiver[ReceiverMessageT_co]:
+        """Take only the messages that fulfill a predicate.
+
+        The returned receiver will only receive messages that fulfill the predicate
+        (evaluates to `True`), and will drop messages that don't.
+
+        Tip:
+            The returned receiver type won't have all the methods of the original
+            receiver. If you need to access methods of the original receiver that are
+            not part of the `Receiver` interface you should save a reference to the
+            original receiver and use that instead.
+
+        Args:
+            predicate: The predicate to be applied on incoming messages to
+                determine if they should be taken.
+
+        Returns:
+            A new receiver that only receives messages that fulfill the predicate.
+        """
+        ...  # pylint: disable=unnecessary-ellipsis
+
+    def take_while(
+        self,
+        predicate: (
+            Callable[[ReceiverMessageT_co], bool]
+            | Callable[[ReceiverMessageT_co], TypeGuard[FilteredMessageT_co]]
+        ),
+        /,
+    ) -> Receiver[ReceiverMessageT_co] | Receiver[FilteredMessageT_co]:
+        """Take only the messages that fulfill a predicate.
+
+        The returned receiver will only receive messages that fulfill the predicate
+        (evaluates to `True`), and will drop messages that don't.
+
+        Note:
+            You can pass a [type guard][typing.TypeGuard] as the predicate to narrow the
+            type of the received messages.
+
+        Tip:
+            The returned receiver type won't have all the methods of the original
+            receiver. If you need to access methods of the original receiver that are
+            not part of the `Receiver` interface you should save a reference to the
+            original receiver and use that instead.
+
+        Args:
+            predicate: The predicate to be applied on incoming messages to
+                determine if they should be taken.
+
+        Returns:
+            A new receiver that only receives messages that fulfill the predicate.
+        """
+        return _Filter(receiver=self, filter_function=predicate)
+
+    def drop_while(
+        self,
+        predicate: Callable[[ReceiverMessageT_co], bool],
+        /,
+    ) -> Receiver[ReceiverMessageT_co] | Receiver[ReceiverMessageT_co]:
+        """Drop the messages that fulfill a predicate.
+
+        The returned receiver will drop messages that fulfill the predicate
+        (evaluates to `True`), and receive messages that don't.
+
+        Tip:
+            If you need to narrow the type of the received messages, you can use the
+            [`take_while()`][frequenz.channels.Receiver.take_while] method instead.
+
+        Tip:
+            The returned receiver type won't have all the methods of the original
+            receiver. If you need to access methods of the original receiver that are
+            not part of the `Receiver` interface you should save a reference to the
+            original receiver and use that instead.
+
+        Args:
+            predicate: The predicate to be applied on incoming messages to
+                determine if they should be dropped.
+
+        Returns:
+            A new receiver that only receives messages that don't fulfill the predicate.
+        """
+        return _Filter(receiver=self, filter_function=predicate, negate=True)
+
     def triggered(
         self, selected: Selected[Any]
     ) -> TypeGuard[Selected[ReceiverMessageT_co]]:
@@ -492,12 +620,14 @@ def __init__(
         *,
         receiver: Receiver[ReceiverMessageT_co],
         filter_function: Callable[[ReceiverMessageT_co], bool],
+        negate: bool = False,
     ) -> None:
         """Initialize this receiver filter.
 
         Args:
             receiver: The input receiver.
             filter_function: The function to apply on the input data.
+            negate: Whether to negate the filter function.
         """
         self._receiver: Receiver[ReceiverMessageT_co] = receiver
         """The input receiver."""
@@ -507,6 +637,8 @@ def __init__(
 
         self._next_message: ReceiverMessageT_co | _Sentinel = _SENTINEL
 
+        self._negate: bool = negate
+
         self._recv_closed = False
 
     async def ready(self) -> bool:
@@ -522,7 +654,10 @@ async def ready(self) -> bool:
         """
         while await self._receiver.ready():
             message = self._receiver.consume()
-            if self._filter_function(message):
+            result = self._filter_function(message)
+            if self._negate:
+                result = not result
+            if result:
                 self._next_message = message
                 return True
         self._recv_closed = True

tests/test_receiver.py

@@ -0,0 +1,87 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Tests for the Receiver class."""
+
+import asyncio
+from collections.abc import Sequence
+from typing import TypeGuard, assert_type
+
+import pytest
+from typing_extensions import override
+
+from frequenz.channels import Receiver, ReceiverError, ReceiverStoppedError
+
+
+class _MockReceiver(Receiver[int | str]):
+    def __init__(self, messages: Sequence[int | str] = ()) -> None:
+        self.messages = list(messages)
+        self.stopped = False
+
+    @override
+    async def ready(self) -> bool:
+        """Return True if there are messages to consume or the receiver is stopped."""
+        if self.stopped:
+            return False
+        if self.messages:
+            await asyncio.sleep(0)
+            return True
+        return False
+
+    @override
+    def consume(self) -> int | str:
+        """Return the next message."""
+        if self.stopped or not self.messages:
+            raise ReceiverStoppedError(self)
+        return self.messages.pop(0)
+
+    def stop(self) -> None:
+        """Stop the receiver."""
+        self.stopped = True
+
+
+async def test_receiver_take_while() -> None:
+    """Test the take_while method."""
+    receiver = _MockReceiver([1, 2, 3, 4, 5])
+
+    filtered_receiver = receiver.take_while(lambda x: x % 2 == 0)
+    async with asyncio.timeout(1):
+        assert await filtered_receiver.receive() == 2
+        assert await filtered_receiver.receive() == 4
+
+        with pytest.raises(ReceiverStoppedError):
+            await filtered_receiver.receive()
+
+
+async def test_receiver_take_type_guard() -> None:
+    """Test the take_while method using a TypeGuard."""
+    receiver = _MockReceiver([1, "1", 2, "2", 3, "3", 4, 5])
+
+    def is_even(x: int | str) -> TypeGuard[int]:
+        if not isinstance(x, int):
+            return False
+        return x % 2 == 0
+
+    filtered_receiver = receiver.take_while(is_even)
+    async with asyncio.timeout(1):
+        received = await filtered_receiver.receive()
+        assert_type(received, int)
+        assert received == 2
+        assert await filtered_receiver.receive() == 4
+
+        with pytest.raises(ReceiverStoppedError):
+            await filtered_receiver.receive()
+
+
+async def test_receiver_drop_while() -> None:
+    """Test the drop_while method."""
+    receiver = _MockReceiver([1, 2, 3, 4, 5])
+
+    filtered_receiver = receiver.drop_while(lambda x: x % 2 == 0)
+    async with asyncio.timeout(1):
+        assert await filtered_receiver.receive() == 1
+        assert await filtered_receiver.receive() == 3
+        assert await filtered_receiver.receive() == 5
+
+        with pytest.raises(ReceiverStoppedError):
+            await filtered_receiver.receive()