Support broadening or narrowing of types in Receivers and Senders (#262)

This allows Receivers to reduce message types to be broader than actual,
and senders to accept narrower types than required.

```python
    @DataClass
    class Broader:
        """A broad class."""

        value: int

    class Actual(Broader):
        """Actual class."""

    class Narrower(Actual):
        """A narrower class."""

    chan = Broadcast[Actual](name="input-chan")

    sender: Sender[Narrower] = chan.new_sender()
    receiver: Receiver[Broader] = chan.new_receiver()
```

Author: shsms

RELEASE_NOTES.md

@@ -147,6 +147,12 @@
 
   - `map()`: The returned map object now has a more useful implementation of `__str__ and `__repr__`.
 
+## Improvements
+
+* `Receiver`: Use a covariant generic type, which allows the generic type to be broader than the actual type.
+
+* `Sender`: Use a contravariant generic type, which allows the generic type to be narrower than the required type.
+
 ## Bug Fixes
 
 * `Timer`: Fix bug that was causing calls to `reset()` to not reset the timer, if the timer was already being awaited.

src/frequenz/channels/_receiver.py

@@ -140,14 +140,14 @@
 
 from ._exceptions import Error
 
-_T = TypeVar("_T")
-_U = TypeVar("_U")
+_T_co = TypeVar("_T_co", covariant=True)
+_U_co = TypeVar("_U_co", covariant=True)
 
 
-class Receiver(ABC, Generic[_T]):
+class Receiver(ABC, Generic[_T_co]):
     """An endpoint to receive messages."""
 
-    async def __anext__(self) -> _T:
+    async def __anext__(self) -> _T_co:
         """Await the next value in the async iteration over received values.
 
         Returns:
@@ -177,7 +177,7 @@ async def ready(self) -> bool:
         """
 
     @abstractmethod
-    def consume(self) -> _T:
+    def consume(self) -> _T_co:
         """Return the latest value once `ready()` is complete.
 
         `ready()` must be called before each call to `consume()`.
@@ -198,7 +198,7 @@ def __aiter__(self) -> Self:
         """
         return self
 
-    async def receive(self) -> _T:
+    async def receive(self) -> _T_co:
         """Receive a message from the channel.
 
         Returns:
@@ -225,7 +225,7 @@ async def receive(self) -> _T:
             raise ReceiverStoppedError(self) from exc
         return received
 
-    def map(self, call: Callable[[_T], _U]) -> Receiver[_U]:
+    def map(self, call: Callable[[_T_co], _U_co]) -> Receiver[_U_co]:
         """Return a receiver with `call` applied on incoming messages.
 
         Args:
@@ -237,13 +237,13 @@ def map(self, call: Callable[[_T], _U]) -> Receiver[_U]:
         return _Map(self, call)
 
 
-class ReceiverError(Error, Generic[_T]):
+class ReceiverError(Error, Generic[_T_co]):
     """An error produced in a [Receiver][frequenz.channels.Receiver].
 
     All exceptions generated by receivers inherit from this exception.
     """
 
-    def __init__(self, message: str, receiver: Receiver[_T]):
+    def __init__(self, message: str, receiver: Receiver[_T_co]):
         """Create an instance.
 
         Args:
@@ -252,14 +252,14 @@ def __init__(self, message: str, receiver: Receiver[_T]):
                 error happened.
         """
         super().__init__(message)
-        self.receiver: Receiver[_T] = receiver
+        self.receiver: Receiver[_T_co] = receiver
         """The receiver where the error happened."""
 
 
-class ReceiverStoppedError(ReceiverError[_T]):
+class ReceiverStoppedError(ReceiverError[_T_co]):
     """The [Receiver][frequenz.channels.Receiver] stopped producing messages."""
 
-    def __init__(self, receiver: Receiver[_T]):
+    def __init__(self, receiver: Receiver[_T_co]):
         """Create an instance.
 
         Args:
@@ -269,7 +269,7 @@ def __init__(self, receiver: Receiver[_T]):
         super().__init__(f"Receiver {receiver} was stopped", receiver)
 
 
-class _Map(Receiver[_U], Generic[_T, _U]):
+class _Map(Receiver[_U_co], Generic[_T_co, _U_co]):
     """Apply a transform function on a channel receiver.
 
     Has two generic types:
@@ -278,17 +278,19 @@ class _Map(Receiver[_U], Generic[_T, _U]):
     - The output type: return type of the transform method.
     """
 
-    def __init__(self, receiver: Receiver[_T], transform: Callable[[_T], _U]) -> None:
+    def __init__(
+        self, receiver: Receiver[_T_co], transform: Callable[[_T_co], _U_co]
+    ) -> None:
         """Create a `Transform` instance.
 
         Args:
             receiver: The input receiver.
             transform: The function to run on the input data.
         """
-        self._receiver: Receiver[_T] = receiver
+        self._receiver: Receiver[_T_co] = receiver
         """The input receiver."""
 
-        self._transform: Callable[[_T], _U] = transform
+        self._transform: Callable[[_T_co], _U_co] = transform
         """The function to run on the input data."""
 
     async def ready(self) -> bool:
@@ -306,7 +308,7 @@ async def ready(self) -> bool:
 
     # We need a noqa here because the docs have a Raises section but the code doesn't
     # explicitly raise anything.
-    def consume(self) -> _U:  # noqa: DOC502
+    def consume(self) -> _U_co:  # noqa: DOC502
         """Return a transformed value once `ready()` is complete.
 
         Returns:

src/frequenz/channels/_sender.py

@@ -54,14 +54,14 @@
 
 from ._exceptions import Error
 
-_T = TypeVar("_T")
+_T_contra = TypeVar("_T_contra", contravariant=True)
 
 
-class Sender(ABC, Generic[_T]):
+class Sender(ABC, Generic[_T_contra]):
     """An endpoint to sends messages."""
 
     @abstractmethod
-    async def send(self, msg: _T) -> None:
+    async def send(self, msg: _T_contra) -> None:
         """Send a message to the channel.
 
         Args:
@@ -72,13 +72,13 @@ async def send(self, msg: _T) -> None:
         """
 
 
-class SenderError(Error, Generic[_T]):
+class SenderError(Error, Generic[_T_contra]):
     """An error produced in a [Sender][frequenz.channels.Sender].
 
     All exceptions generated by senders inherit from this exception.
     """
 
-    def __init__(self, message: str, sender: Sender[_T]):
+    def __init__(self, message: str, sender: Sender[_T_contra]):
         """Create an instance.
 
         Args:
@@ -87,5 +87,5 @@ def __init__(self, message: str, sender: Sender[_T]):
                 happened.
         """
         super().__init__(message)
-        self.sender: Sender[_T] = sender
+        self.sender: Sender[_T_contra] = sender
         """The sender where the error happened."""

tests/test_broadcast.py

@@ -5,6 +5,7 @@
 
 
 import asyncio
+from dataclasses import dataclass
 
 import pytest
 
@@ -252,3 +253,27 @@ async def test_broadcast_receiver_drop() -> None:
 
     assert len(chan._receivers) == 1
     # pylint: enable=protected-access
+
+
+async def test_type_variance() -> None:
+    """Ensure that the type variance of Broadcast is working."""
+
+    @dataclass
+    class Broader:
+        """A broad class."""
+
+        value: int
+
+    class Actual(Broader):
+        """Actual class."""
+
+    class Narrower(Actual):
+        """A narrower class."""
+
+    chan = Broadcast[Actual](name="input-chan")
+
+    sender: Sender[Narrower] = chan.new_sender()
+    receiver: Receiver[Broader] = chan.new_receiver()
+
+    await sender.send(Narrower(10))
+    assert (await receiver.receive()).value == 10