Add generic typing support for MemorySendChannel and MemoryReceiveCha…

…nnel

trio/_channel.py

@@ -1,18 +1,36 @@
 from collections import deque, OrderedDict
 from math import inf
 
+from types import TracebackType
+from typing import Generic, TypeVar, Any, Tuple, Optional, Type
+
 import attr
 from outcome import Error, Value
 
 from .abc import SendChannel, ReceiveChannel, Channel
 from ._util import generic_function, NoPublicConstructor
 
 import trio
-from ._core import enable_ki_protection
+from ._core import enable_ki_protection, Task, Abort
+
+# A regular invariant generic type
+T = TypeVar("T")
+
+# The type of object produced by a ReceiveChannel (covariant because
+# ReceiveChannel[Derived] can be passed to someone expecting
+# ReceiveChannel[Base])
+ReceiveType = TypeVar("ReceiveType", covariant=True)
+
+# The type of object accepted by a SendChannel (contravariant because
+# SendChannel[Base] can be passed to someone expecting
+# SendChannel[Derived])
+SendType = TypeVar("SendType", contravariant=True)
 
 
 @generic_function
-def open_memory_channel(max_buffer_size):
+def open_memory_channel(
+ max_buffer_size: int,
+) -> Tuple[SendChannel[T], ReceiveChannel[T]]:
  """Open a channel for passing objects between tasks within a process.
 
  Memory channels are lightweight, cheap to allocate, and entirely
@@ -68,7 +86,7 @@ def open_memory_channel(max_buffer_size):
  raise TypeError("max_buffer_size must be an integer or math.inf")
  if max_buffer_size < 0:
  raise ValueError("max_buffer_size must be >= 0")
- state = MemoryChannelState(max_buffer_size)
+ state: MemoryChannelState[T] = MemoryChannelState(max_buffer_size)
  return (
  MemorySendChannel._create(state),
  MemoryReceiveChannel._create(state),
@@ -77,27 +95,27 @@ def open_memory_channel(max_buffer_size):
 
 @attr.s(frozen=True, slots=True)
 class MemoryChannelStats:
- current_buffer_used = attr.ib()
- max_buffer_size = attr.ib()
- open_send_channels = attr.ib()
- open_receive_channels = attr.ib()
- tasks_waiting_send = attr.ib()
- tasks_waiting_receive = attr.ib()
+ current_buffer_used: int = attr.ib()
+ max_buffer_size: int = attr.ib()
+ open_send_channels: int = attr.ib()
+ open_receive_channels: int = attr.ib()
+ tasks_waiting_send: int = attr.ib()
+ tasks_waiting_receive: int = attr.ib()
 
 
 @attr.s(slots=True)
-class MemoryChannelState:
- max_buffer_size = attr.ib()
- data = attr.ib(factory=deque)
+class MemoryChannelState(Generic[T]):
+ max_buffer_size: int = attr.ib()
+ data: deque[T] = attr.ib(factory=deque)
  # Counts of open endpoints using this state
- open_send_channels = attr.ib(default=0)
- open_receive_channels = attr.ib(default=0)
+ open_send_channels: int = attr.ib(default=0)
+ open_receive_channels: int = attr.ib(default=0)
  # {task: value}
- send_tasks = attr.ib(factory=OrderedDict)
+ send_tasks: OrderedDict[Task, T] = attr.ib(factory=OrderedDict)
  # {task: None}
- receive_tasks = attr.ib(factory=OrderedDict)
+ receive_tasks: OrderedDict[Task, None] = attr.ib(factory=OrderedDict)
 
- def statistics(self):
+ def statistics(self) -> MemoryChannelStats:
  return MemoryChannelStats(
  current_buffer_used=len(self.data),
  max_buffer_size=self.max_buffer_size,
@@ -109,28 +127,28 @@ def statistics(self):
 
 
 @attr.s(eq=False, repr=False)
-class MemorySendChannel(SendChannel, metaclass=NoPublicConstructor):
- _state = attr.ib()
- _closed = attr.ib(default=False)
+class MemorySendChannel(SendChannel, Generic[SendType], metaclass=NoPublicConstructor):
+ _state: MemoryChannelState[SendType] = attr.ib()
+ _closed: bool = attr.ib(default=False)
  # This is just the tasks waiting on *this* object. As compared to
  # self._state.send_tasks, which includes tasks from this object and
  # all clones.
- _tasks = attr.ib(factory=set)
+ _tasks: set[Task] = attr.ib(factory=set)
 
- def __attrs_post_init__(self):
+ def __attrs_post_init__(self) -> None:
  self._state.open_send_channels += 1
 
- def __repr__(self):
+ def __repr__(self) -> str:
  return "<send channel at {:#x}, using buffer at {:#x}>".format(
  id(self), id(self._state)
  )
 
- def statistics(self):
+ def statistics(self) -> MemoryChannelStats:
  # XX should we also report statistics specific to this object?
  return self._state.statistics()
 
  @enable_ki_protection
- def send_nowait(self, value):
+ def send_nowait(self, value: SendType) -> None:
  """Like `~trio.abc.SendChannel.send`, but if the channel's buffer is
  full, raises `WouldBlock` instead of blocking.
 
@@ -150,7 +168,7 @@ def send_nowait(self, value):
  raise trio.WouldBlock
 
  @enable_ki_protection
- async def send(self, value):
+ async def send(self, value: SendType) -> None:
  """See `SendChannel.send <trio.abc.SendChannel.send>`.
 
  Memory channels allow multiple tasks to call `send` at the same time.
@@ -170,15 +188,16 @@ async def send(self, value):
  self._state.send_tasks[task] = value
  task.custom_sleep_data = self
 
- def abort_fn(_):
+ def abort_fn(_) -> Abort:
  self._tasks.remove(task)
  del self._state.send_tasks[task]
  return trio.lowlevel.Abort.SUCCEEDED
 
  await trio.lowlevel.wait_task_rescheduled(abort_fn)
 
+ # Return type must be stringified, use a TypeVar, or (py311+) use typing.Self
  @enable_ki_protection
- def clone(self):
+ def clone(self) -> "MemorySendChannel[SendType]":
  """Clone this send channel object.
 
  This returns a new `MemorySendChannel` object, which acts as a
@@ -206,14 +225,19 @@ def clone(self):
  raise trio.ClosedResourceError
  return MemorySendChannel._create(self._state)
 
- def __enter__(self):
+ def __enter__(self) -> "MemorySendChannel[SendType]":
  return self
 
- def __exit__(self, exc_type, exc_val, exc_tb):
+ def __exit__(
+ self,
+ exc_type: Optional[Type[BaseException]],
+ exc_val: Optional[BaseException],
+ exc_tb: Optional[TracebackType],
+ ) -> None:
  self.close()
 
  @enable_ki_protection
- def close(self):
+ def close(self) -> None:
  """Close this send channel object synchronously.
 
  All channel objects have an asynchronous `~.AsyncResource.aclose` method.
@@ -241,30 +265,32 @@ def close(self):
  self._state.receive_tasks.clear()
 
  @enable_ki_protection
- async def aclose(self):
+ async def aclose(self) -> None:
  self.close()
  await trio.lowlevel.checkpoint()
 
 
 @attr.s(eq=False, repr=False)
-class MemoryReceiveChannel(ReceiveChannel, metaclass=NoPublicConstructor):
- _state = attr.ib()
- _closed = attr.ib(default=False)
- _tasks = attr.ib(factory=set)
-
- def __attrs_post_init__(self):
+class MemoryReceiveChannel(
+ ReceiveChannel, Generic[ReceiveType], metaclass=NoPublicConstructor
+):
+ _state: MemoryChannelState[ReceiveType] = attr.ib()
+ _closed: bool = attr.ib(default=False)
+ _tasks: set[trio._core._run.Task] = attr.ib(factory=set)
+
+ def __attrs_post_init__(self) -> None:
  self._state.open_receive_channels += 1
 
- def statistics(self):
+ def statistics(self) -> MemoryChannelStats:
  return self._state.statistics()
 
- def __repr__(self):
+ def __repr__(self) -> str:
  return "<receive channel at {:#x}, using buffer at {:#x}>".format(
  id(self), id(self._state)
  )
 
  @enable_ki_protection
- def receive_nowait(self):
+ def receive_nowait(self) -> ReceiveType:
  """Like `~trio.abc.ReceiveChannel.receive`, but if there's nothing
  ready to receive, raises `WouldBlock` instead of blocking.
 
@@ -284,7 +310,7 @@ def receive_nowait(self):
  raise trio.WouldBlock
 
  @enable_ki_protection
- async def receive(self):
+ async def receive(self) -> ReceiveType:
  """See `ReceiveChannel.receive <trio.abc.ReceiveChannel.receive>`.
 
  Memory channels allow multiple tasks to call `receive` at the same
@@ -306,15 +332,15 @@ async def receive(self):
  self._state.receive_tasks[task] = None
  task.custom_sleep_data = self
 
- def abort_fn(_):
+ def abort_fn(_) -> Abort:
  self._tasks.remove(task)
  del self._state.receive_tasks[task]
  return trio.lowlevel.Abort.SUCCEEDED
 
- return await trio.lowlevel.wait_task_rescheduled(abort_fn)
+ return await trio.lowlevel.wait_task_rescheduled(abort_fn) # type: ignore
 
  @enable_ki_protection
- def clone(self):
+ def clone(self) -> "MemoryReceiveChannel[ReceiveType]":
  """Clone this receive channel object.
 
  This returns a new `MemoryReceiveChannel` object, which acts as a
@@ -345,14 +371,19 @@ def clone(self):
  raise trio.ClosedResourceError
  return MemoryReceiveChannel._create(self._state)
 
- def __enter__(self):
+ def __enter__(self) -> "MemoryReceiveChannel[ReceiveType]":
  return self
 
- def __exit__(self, exc_type, exc_val, exc_tb):
+ def __exit__(
+ self,
+ exc_type: Optional[Type[BaseException]],
+ exc_val: Optional[BaseException],
+ exc_tb: Optional[TracebackType],
+ ) -> None:
  self.close()
 
  @enable_ki_protection
- def close(self):
+ def close(self) -> None:
  """Close this receive channel object synchronously.
 
  All channel objects have an asynchronous `~.AsyncResource.aclose` method.
@@ -381,6 +412,6 @@ def close(self):
  self._state.data.clear()
 
  @enable_ki_protection
- async def aclose(self):
+ async def aclose(self) -> None:
  self.close()
  await trio.lowlevel.checkpoint()