Use new type hinting for trio channel

This commit is contained in:
mhchia 2020-01-26 23:56:19 +08:00
parent 42bc4d5d06
commit e3a1dd62e4
No known key found for this signature in database
GPG Key ID: 389EFBEA1362589A
2 changed files with 5 additions and 18 deletions

View File

@ -71,7 +71,6 @@ class Pubsub(IPubsub, Service):
seen_messages: LRU seen_messages: LRU
# TODO: Implement `trio.abc.Channel`?
subscribed_topics_send: Dict[str, "trio.MemorySendChannel[rpc_pb2.Message]"] subscribed_topics_send: Dict[str, "trio.MemorySendChannel[rpc_pb2.Message]"]
subscribed_topics_receive: Dict[str, "TrioSubscriptionAPI"] subscribed_topics_receive: Dict[str, "TrioSubscriptionAPI"]
@ -112,12 +111,8 @@ class Pubsub(IPubsub, Service):
# Attach this new Pubsub object to the router # Attach this new Pubsub object to the router
self.router.attach(self) self.router.attach(self)
peer_channels: Tuple[ peer_channels = trio.open_memory_channel[ID](0)
"trio.MemorySendChannel[ID]", "trio.MemoryReceiveChannel[ID]" dead_peer_channels = trio.open_memory_channel[ID](0)
] = trio.open_memory_channel(0)
dead_peer_channels: Tuple[
"trio.MemorySendChannel[ID]", "trio.MemoryReceiveChannel[ID]"
] = trio.open_memory_channel(0)
# Only keep the receive channels in `Pubsub`. # Only keep the receive channels in `Pubsub`.
# Therefore, we can only close from the receive side. # Therefore, we can only close from the receive side.
self.peer_receive_channel = peer_channels[1] self.peer_receive_channel = peer_channels[1]
@ -404,10 +399,7 @@ class Pubsub(IPubsub, Service):
if topic_id in self.topic_ids: if topic_id in self.topic_ids:
return self.subscribed_topics_receive[topic_id] return self.subscribed_topics_receive[topic_id]
channels: Tuple[ channels = trio.open_memory_channel[rpc_pb2.Message](math.inf)
"trio.MemorySendChannel[rpc_pb2.Message]",
"trio.MemoryReceiveChannel[rpc_pb2.Message]",
] = trio.open_memory_channel(math.inf)
send_channel, receive_channel = channels send_channel, receive_channel = channels
subscription = TrioSubscriptionAPI(receive_channel) subscription = TrioSubscriptionAPI(receive_channel)
self.subscribed_topics_send[topic_id] = send_channel self.subscribed_topics_send[topic_id] = send_channel

View File

@ -66,10 +66,7 @@ class Mplex(IMuxedConn):
self.streams = {} self.streams = {}
self.streams_lock = trio.Lock() self.streams_lock = trio.Lock()
self.streams_msg_channels = {} self.streams_msg_channels = {}
channels: Tuple[ channels = trio.open_memory_channel[IMuxedStream](math.inf)
"trio.MemorySendChannel[IMuxedStream]",
"trio.MemoryReceiveChannel[IMuxedStream]",
] = trio.open_memory_channel(math.inf)
self.new_stream_send_channel, self.new_stream_receive_channel = channels self.new_stream_send_channel, self.new_stream_receive_channel = channels
self.event_shutting_down = trio.Event() self.event_shutting_down = trio.Event()
self.event_closed = trio.Event() self.event_closed = trio.Event()
@ -114,9 +111,7 @@ class Mplex(IMuxedConn):
async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream: async def _initialize_stream(self, stream_id: StreamID, name: str) -> MplexStream:
# Use an unbounded buffer, to avoid `handle_incoming` being blocked when doing # Use an unbounded buffer, to avoid `handle_incoming` being blocked when doing
# `send_channel.send`. # `send_channel.send`.
channels: Tuple[ channels = trio.open_memory_channel[bytes](math.inf)
"trio.MemorySendChannel[bytes]", "trio.MemoryReceiveChannel[bytes]"
] = trio.open_memory_channel(math.inf)
stream = MplexStream(name, stream_id, self, channels[1]) stream = MplexStream(name, stream_id, self, channels[1])
async with self.streams_lock: async with self.streams_lock:
self.streams[stream_id] = stream self.streams[stream_id] = stream