Use the new public type variables

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

Author: llucax

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/_receiver.py

@@ -136,18 +136,16 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from typing import Generic, Self, TypeVar
+from typing import Generic, Self
 
 from ._exceptions import Error
+from ._generic import MappedMessageT_co, ReceiverMessageT_co
 
-_T_co = TypeVar("_T_co", covariant=True)
-_U_co = TypeVar("_U_co", covariant=True)
 
-
-class Receiver(ABC, Generic[_T_co]):
+class Receiver(ABC, Generic[ReceiverMessageT_co]):
     """An endpoint to receive messages."""
 
-    async def __anext__(self) -> _T_co:
+    async def __anext__(self) -> ReceiverMessageT_co:
         """Await the next message in the async iteration over received messages.
 
         Returns:
@@ -177,7 +175,7 @@ async def ready(self) -> bool:
         """
 
     @abstractmethod
-    def consume(self) -> _T_co:
+    def consume(self) -> ReceiverMessageT_co:
         """Return the latest message once `ready()` is complete.
 
         `ready()` must be called before each call to `consume()`.
@@ -198,7 +196,7 @@ def __aiter__(self) -> Self:
         """
         return self
 
-    async def receive(self) -> _T_co:
+    async def receive(self) -> ReceiverMessageT_co:
         """Receive a message.
 
         Returns:
@@ -225,7 +223,9 @@ async def receive(self) -> _T_co:
             raise ReceiverStoppedError(self) from exc
         return received
 
-    def map(self, mapping_function: Callable[[_T_co], _U_co], /) -> Receiver[_U_co]:
+    def map(
+        self, mapping_function: Callable[[ReceiverMessageT_co], MappedMessageT_co], /
+    ) -> Receiver[MappedMessageT_co]:
         """Apply a mapping function on the received message.
 
         Tip:
@@ -243,13 +243,13 @@ def map(self, mapping_function: Callable[[_T_co], _U_co], /) -> Receiver[_U_co]:
         return _Mapper(receiver=self, mapping_function=mapping_function)
 
 
-class ReceiverError(Error, Generic[_T_co]):
+class ReceiverError(Error, Generic[ReceiverMessageT_co]):
     """An error that originated in a [Receiver][frequenz.channels.Receiver].
 
     All exceptions generated by receivers inherit from this exception.
     """
 
-    def __init__(self, message: str, receiver: Receiver[_T_co]):
+    def __init__(self, message: str, receiver: Receiver[ReceiverMessageT_co]):
         """Initialize this error.
 
         Args:
@@ -258,14 +258,14 @@ def __init__(self, message: str, receiver: Receiver[_T_co]):
                 error happened.
         """
         super().__init__(message)
-        self.receiver: Receiver[_T_co] = receiver
+        self.receiver: Receiver[ReceiverMessageT_co] = receiver
         """The receiver where the error happened."""
 
 
-class ReceiverStoppedError(ReceiverError[_T_co]):
+class ReceiverStoppedError(ReceiverError[ReceiverMessageT_co]):
     """A stopped [`Receiver`][frequenz.channels.Receiver] was used."""
 
-    def __init__(self, receiver: Receiver[_T_co]):
+    def __init__(self, receiver: Receiver[ReceiverMessageT_co]):
         """Initialize this error.
 
         Args:
@@ -275,7 +275,9 @@ def __init__(self, receiver: Receiver[_T_co]):
         super().__init__(f"Receiver {receiver} was stopped", receiver)
 
 
-class _Mapper(Receiver[_U_co], Generic[_T_co, _U_co]):
+class _Mapper(
+    Receiver[MappedMessageT_co], Generic[ReceiverMessageT_co, MappedMessageT_co]
+):
     """Apply a transform function on a channel receiver.
 
     Has two generic types:
@@ -285,18 +287,23 @@ class _Mapper(Receiver[_U_co], Generic[_T_co, _U_co]):
     """
 
     def __init__(
-        self, *, receiver: Receiver[_T_co], mapping_function: Callable[[_T_co], _U_co]
+        self,
+        *,
+        receiver: Receiver[ReceiverMessageT_co],
+        mapping_function: Callable[[ReceiverMessageT_co], MappedMessageT_co],
     ) -> None:
         """Initialize this receiver mapper.
 
         Args:
             receiver: The input receiver.
             mapping_function: The function to apply on the input data.
         """
-        self._receiver: Receiver[_T_co] = receiver
+        self._receiver: Receiver[ReceiverMessageT_co] = receiver
         """The input receiver."""
 
-        self._mapping_function: Callable[[_T_co], _U_co] = mapping_function
+        self._mapping_function: Callable[[ReceiverMessageT_co], MappedMessageT_co] = (
+            mapping_function
+        )
         """The function to apply on the input data."""
 
     async def ready(self) -> bool:
@@ -314,7 +321,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_co:  # noqa: DOC502
+    def consume(self) -> MappedMessageT_co:  # noqa: DOC502
         """Return a transformed message once `ready()` is complete.
 
         Returns:

src/frequenz/channels/_sender.py

@@ -50,18 +50,17 @@
 """
 
 from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
+from typing import Generic
 
 from ._exceptions import Error
+from ._generic import SenderMessageT_contra
 
-_T_contra = TypeVar("_T_contra", contravariant=True)
 
-
-class Sender(ABC, Generic[_T_contra]):
+class Sender(ABC, Generic[SenderMessageT_contra]):
     """An endpoint to sends messages."""
 
     @abstractmethod
-    async def send(self, message: _T_contra, /) -> None:
+    async def send(self, message: SenderMessageT_contra, /) -> None:
         """Send a message.
 
         Args:
@@ -72,13 +71,13 @@ async def send(self, message: _T_contra, /) -> None:
         """
 
 
-class SenderError(Error, Generic[_T_contra]):
+class SenderError(Error, Generic[SenderMessageT_contra]):
     """An error that originated in a [Sender][frequenz.channels.Sender].
 
     All exceptions generated by senders inherit from this exception.
     """
 
-    def __init__(self, message: str, sender: Sender[_T_contra]):
+    def __init__(self, message: str, sender: Sender[SenderMessageT_contra]):
         """Initialize this error.
 
         Args:
@@ -87,5 +86,5 @@ def __init__(self, message: str, sender: Sender[_T_contra]):
                 happened.
         """
         super().__init__(message)
-        self.sender: Sender[_T_contra] = sender
+        self.sender: Sender[SenderMessageT_contra] = sender
         """The sender where the error happened."""