Use covariant generic types for Receivers

This allows broadening of the message type, such that receivers that
stream narrow types can be typed as receivers streaming broader types.

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

Author: shsms

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: