diff --git a/ssb/muxrpc.py b/ssb/muxrpc.py index 4832695..b0abf65 100644 --- a/ssb/muxrpc.py +++ b/ssb/muxrpc.py @@ -134,15 +134,22 @@ class MuxRPCAPI: self.handlers = {} self.connection = None - async def __await__(self): - async for req_message in self.connection: - if req_message is None: - return + def __aiter__(self): + return self - body = req_message.body + async def __anext__(self): + req_message = await self.connection.__anext__() - if isinstance(body, dict) and body.get("name"): - self.process(self.connection, MuxRPCRequest.from_message(req_message)) + if req_message is None: + raise StopAsyncIteration() + + body = req_message.body + + if isinstance(body, dict) and body.get("name"): + self.process(self.connection, MuxRPCRequest.from_message(req_message)) + + def __await__(self): + yield from self.__anext__().__await__() def add_connection(self, connection): """Set the packet stream connection of this RPC API"""