From f2a54b5ce6e3b2919a06fc690e1466252c25afc0 Mon Sep 17 00:00:00 2001
From: Gergely Polonkai <gergely@polonkai.eu>
Date: Thu, 16 Nov 2023 05:40:08 +0100
Subject: [PATCH] chore: Get rid of all __await__ methods

---
 examples/test_client.py |  5 ++++-
 ssb/muxrpc.py           |  8 ++++++--
 ssb/packet_stream.py    | 16 +++++++---------
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/examples/test_client.py b/examples/test_client.py
index c08907b..536a301 100644
--- a/examples/test_client.py
+++ b/examples/test_client.py
@@ -73,7 +73,10 @@ async def test_client():
         print("> RESPONSE:", msg)
 
     try:
-        print("> RESPONSE:", await api.call("whoami", [], "sync"))
+        response_handler = api.call("whoami", [], "sync")
+        response = await response_handler.get_response()
+
+        print("> RESPONSE:", response)
     except MuxRPCAPIException as e:
         print(e)
 
diff --git a/ssb/muxrpc.py b/ssb/muxrpc.py
index 170cd52..a059196 100644
--- a/ssb/muxrpc.py
+++ b/ssb/muxrpc.py
@@ -47,9 +47,13 @@ class MuxRPCRequestHandler(MuxRPCHandler):
     def __init__(self, ps_handler):
         self.ps_handler = ps_handler
 
-    def __await__(self):
-        msg = yield from self.ps_handler.__await__()
+    async def get_response(self):
+        """Get the response data"""
+
+        msg = await self.ps_handler
+
         self.check_message(msg)
+
         return msg
 
 
diff --git a/ssb/packet_stream.py b/ssb/packet_stream.py
index b51e25c..f61941e 100644
--- a/ssb/packet_stream.py
+++ b/ssb/packet_stream.py
@@ -92,9 +92,12 @@ class PSRequestHandler:
         if not self.event.is_set():
             self.event.set()
 
-    def __await__(self):
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
         # wait until 'process' is called
-        yield from self.event.wait().__await__()  # pylint: disable=no-member
+        await self.event.wait()
 
         return self._msg
 
@@ -178,15 +181,10 @@ class PacketStream:
                 raise StopAsyncIteration()
 
             if msg.req >= 0:
+                logger.info("RECV: %r", msg)
+
                 return msg
 
-    async def __await__(self):
-        async for data in self:
-            logger.info("RECV: %r", data)
-
-            if data is None:
-                return
-
     async def _read(self):
         try:
             header = await self.connection.read()