Improve generics usage (#282)

This PR makes some improvements related to generic types.

It adds public generic type vars so the documentation is rendered more
nicely and we get documentation on the type parameters.

It also extends the use of covariance for other classes or functions
that work like containers that don't change the underlying data, so they
can be used more flexibly.

Author: llucax

RELEASE_NOTES.md

@@ -173,9 +173,11 @@
 
 ## Improvements
 
-* `Receiver`: Use a covariant generic type, which allows the generic type to be broader than the actual type.
+* `Receiver`, `merge`/`Merger`, `Error` and its derived classes now 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.
+* `Sender` now uses a contravariant generic type, which allows the generic type to be narrower than the required type.
+
+*  `ChannelError` is now generic, so when accessing the `channel` attribute, the type of the channel is preserved.
 
 ## Bug Fixes
 

src/frequenz/channels/__init__.py

@@ -77,6 +77,14 @@
 from ._anycast import Anycast
 from ._broadcast import Broadcast
 from ._exceptions import ChannelClosedError, ChannelError, Error
+from ._generic import (
+    ChannelMessageT,
+    ErroredChannelT_co,
+    MappedMessageT_co,
+    ReceiverMessageT_co,
+    SenderMessageT_co,
+    SenderMessageT_contra,
+)
 from ._merge import Merger, merge
 from ._receiver import Receiver, ReceiverError, ReceiverStoppedError
 from ._select import (
@@ -93,15 +101,21 @@
     "Broadcast",
     "ChannelClosedError",
     "ChannelError",
+    "ChannelMessageT",
     "Error",
+    "ErroredChannelT_co",
+    "MappedMessageT_co",
     "Merger",
     "Receiver",
     "ReceiverError",
+    "ReceiverMessageT_co",
     "ReceiverStoppedError",
     "SelectError",
     "Selected",
     "Sender",
     "SenderError",
+    "SenderMessageT_co",
+    "SenderMessageT_contra",
     "UnhandledSelectedError",
     "merge",
     "select",

src/frequenz/channels/_anycast.py

@@ -11,15 +11,14 @@
 from typing import Generic, TypeVar
 
 from ._exceptions import ChannelClosedError
+from ._generic import ChannelMessageT
 from ._receiver import Receiver, ReceiverStoppedError
 from ._sender import Sender, SenderError
 
 _logger = logging.getLogger(__name__)
 
-_T = TypeVar("_T")
-
 
-class Anycast(Generic[_T]):
+class Anycast(Generic[ChannelMessageT]):
     """A channel that delivers each message to exactly one receiver.
 
     # Description
@@ -213,7 +212,7 @@ def __init__(self, *, name: str, limit: int = 10) -> None:
         of the channel.
         """
 
-        self._deque: deque[_T] = deque(maxlen=limit)
+        self._deque: deque[ChannelMessageT] = deque(maxlen=limit)
         """The channel's buffer."""
 
         self._send_cv: Condition = Condition()
@@ -282,11 +281,11 @@ async def close(self) -> None:
         async with self._recv_cv:
             self._recv_cv.notify_all()
 
-    def new_sender(self) -> Sender[_T]:
+    def new_sender(self) -> Sender[ChannelMessageT]:
         """Return a new sender attached to this channel."""
         return _Sender(self)
 
-    def new_receiver(self) -> Receiver[_T]:
+    def new_receiver(self) -> Receiver[ChannelMessageT]:
         """Return a new receiver attached to this channel."""
         return _Receiver(self)
 
@@ -302,6 +301,9 @@ def __repr__(self) -> str:
         )
 
 
+_T = TypeVar("_T")
+
+
 class _Sender(Sender[_T]):
     """A sender to send messages to an Anycast channel.
 

src/frequenz/channels/_broadcast.py

@@ -12,15 +12,14 @@
 from typing import Generic, TypeVar
 
 from ._exceptions import ChannelClosedError
+from ._generic import ChannelMessageT
 from ._receiver import Receiver, ReceiverStoppedError
 from ._sender import Sender, SenderError
 
 _logger = logging.Logger(__name__)
 
-_T = TypeVar("_T")
-
 
-class Broadcast(Generic[_T]):
+class Broadcast(Generic[ChannelMessageT]):
     """A channel that deliver all messages to all receivers.
 
     # Description
@@ -206,13 +205,15 @@ def __init__(self, *, name: str, resend_latest: bool = False) -> None:
         self._recv_cv: Condition = Condition()
         """The condition to wait for data in the channel's buffer."""
 
-        self._receivers: dict[int, weakref.ReferenceType[_Receiver[_T]]] = {}
+        self._receivers: dict[
+            int, weakref.ReferenceType[_Receiver[ChannelMessageT]]
+        ] = {}
         """The receivers attached to the channel, indexed by their hash()."""
 
         self._closed: bool = False
         """Whether the channel is closed."""
 
-        self._latest: _T | None = None
+        self._latest: ChannelMessageT | None = None
         """The latest message sent to the channel."""
 
         self.resend_latest: bool = resend_latest
@@ -261,11 +262,13 @@ async def close(self) -> None:
         async with self._recv_cv:
             self._recv_cv.notify_all()
 
-    def new_sender(self) -> Sender[_T]:
+    def new_sender(self) -> Sender[ChannelMessageT]:
         """Return a new sender attached to this channel."""
         return _Sender(self)
 
-    def new_receiver(self, *, name: str | None = None, limit: int = 50) -> Receiver[_T]:
+    def new_receiver(
+        self, *, name: str | None = None, limit: int = 50
+    ) -> Receiver[ChannelMessageT]:
         """Return a new receiver attached to this channel.
 
         Broadcast receivers have their own buffer, and when messages are not
@@ -279,7 +282,7 @@ def new_receiver(self, *, name: str | None = None, limit: int = 50) -> Receiver[
         Returns:
             A new receiver attached to this channel.
         """
-        recv: _Receiver[_T] = _Receiver(self, name=name, limit=limit)
+        recv: _Receiver[ChannelMessageT] = _Receiver(self, name=name, limit=limit)
         self._receivers[hash(recv)] = weakref.ref(recv)
         if self.resend_latest and self._latest is not None:
             recv.enqueue(self._latest)
@@ -300,6 +303,9 @@ def __repr__(self) -> str:
         )
 
 
+_T = TypeVar("_T")
+
+
 class _Sender(Sender[_T]):
     """A sender to send messages to the broadcast channel.
 

src/frequenz/channels/_exceptions.py

@@ -66,7 +66,9 @@
     ```
 """
 
-from typing import Any
+from typing import Generic
+
+from ._generic import ErroredChannelT_co
 
 
 class Error(RuntimeError):
@@ -84,28 +86,28 @@ def __init__(self, message: str):
         super().__init__(message)
 
 
-class ChannelError(Error):
+class ChannelError(Error, Generic[ErroredChannelT_co]):
     """An error that originated in a channel.
 
     All exceptions generated by channels inherit from this exception.
     """
 
-    def __init__(self, message: str, channel: Any):
+    def __init__(self, message: str, channel: ErroredChannelT_co):
         """Initialize this error.
 
         Args:
             message: The error message.
             channel: The channel where the error happened.
         """
         super().__init__(message)
-        self.channel: Any = channel
+        self.channel: ErroredChannelT_co = channel
         """The channel where the error happened."""
 
 
-class ChannelClosedError(ChannelError):
+class ChannelClosedError(ChannelError[ErroredChannelT_co]):
     """A closed channel was used."""
 
-    def __init__(self, channel: Any):
+    def __init__(self, channel: ErroredChannelT_co):
         """Initialize this error.
 
         Args:

src/frequenz/channels/_generic.py

@@ -0,0 +1,24 @@
+# License: MIT
+# Copyright © 2024 Frequenz Energy-as-a-Service GmbH
+
+"""Generic type variables."""
+
+from typing import TypeVar
+
+ChannelMessageT = TypeVar("ChannelMessageT")
+"""The type of the message that can be sent across a channel."""
+
+ErroredChannelT_co = TypeVar("ErroredChannelT_co", covariant=True)
+"""The type of channel having an error."""
+
+MappedMessageT_co = TypeVar("MappedMessageT_co", covariant=True)
+"""The type of the message received by the receiver after being mapped."""
+
+ReceiverMessageT_co = TypeVar("ReceiverMessageT_co", covariant=True)
+"""The type of the message received by a receiver."""
+
+SenderMessageT_co = TypeVar("SenderMessageT_co", covariant=True)
+"""The type of the message sent by a sender."""
+
+SenderMessageT_contra = TypeVar("SenderMessageT_contra", contravariant=True)
+"""The type of the message sent by a sender."""

src/frequenz/channels/_merge.py

@@ -52,14 +52,13 @@
 import asyncio
 import itertools
 from collections import deque
-from typing import Any, TypeVar
+from typing import Any
 
+from ._generic import ReceiverMessageT_co
 from ._receiver import Receiver, ReceiverStoppedError
 
-_T = TypeVar("_T")
 
-
-def merge(*receivers: Receiver[_T]) -> Merger[_T]:
+def merge(*receivers: Receiver[ReceiverMessageT_co]) -> Merger[ReceiverMessageT_co]:
     """Merge messages coming from multiple receivers into a single stream.
 
     Example:
@@ -95,31 +94,33 @@ def merge(*receivers: Receiver[_T]) -> Merger[_T]:
     return Merger(*receivers, name="merge")
 
 
-class Merger(Receiver[_T]):
+class Merger(Receiver[ReceiverMessageT_co]):
     """A receiver that merges messages coming from multiple receivers into a single stream.
 
     Tip:
         Please consider using the more idiomatic [`merge()`][frequenz.channels.merge]
         function instead of creating a `Merger` instance directly.
     """
 
-    def __init__(self, *receivers: Receiver[_T], name: str | None) -> None:
+    def __init__(
+        self, *receivers: Receiver[ReceiverMessageT_co], name: str | None
+    ) -> None:
         """Initialize this merger.
 
         Args:
             *receivers: The receivers to merge.
             name: The name of the receiver. Used to create the string representation
                 of the receiver.
         """
-        self._receivers: dict[str, Receiver[_T]] = {
+        self._receivers: dict[str, Receiver[ReceiverMessageT_co]] = {
             str(id): recv for id, recv in enumerate(receivers)
         }
         self._name: str = name if name is not None else type(self).__name__
         self._pending: set[asyncio.Task[Any]] = {
             asyncio.create_task(anext(recv), name=name)
             for name, recv in self._receivers.items()
         }
-        self._results: deque[_T] = deque(maxlen=len(self._receivers))
+        self._results: deque[ReceiverMessageT_co] = deque(maxlen=len(self._receivers))
 
     def __del__(self) -> None:
         """Finalize this merger."""
@@ -170,7 +171,7 @@ async def ready(self) -> bool:
                     asyncio.create_task(anext(self._receivers[name]), name=name)
                 )
 
-    def consume(self) -> _T:
+    def consume(self) -> ReceiverMessageT_co:
         """Return the latest message once `ready` is complete.
 
         Returns: