From a8c3374214aed00a131b7a6fe3c5ee158d0ceba4 Mon Sep 17 00:00:00 2001
From: Arama <arama.dev@outlook.com>
Date: Thu, 19 Dec 2024 19:29:38 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E9=87=8D=E6=96=B0=E6=B7=BB=E5=8A=A0=5F?=
 =?UTF-8?q?=5Fanext=5F=5F=E5=88=B0Connection?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/steam/core/connection.py | 42 ++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/src/steam/core/connection.py b/src/steam/core/connection.py
index 40639ab..3562284 100644
--- a/src/steam/core/connection.py
+++ b/src/steam/core/connection.py
@@ -27,7 +27,7 @@ class Connection:
 
     @property
     def local_address(self):
-        return self._stream_writer.get_extra_info('sockname')[0]
+        raise NotImplementedError
 
     async def connect(self, server_addr):
         raise NotImplementedError
@@ -48,7 +48,7 @@ async def disconnect(self):
         self._readbuf = b''
         self.send_queue._queue.clear()
         self.recv_queue._queue.clear()
-        await self.recv_queue.put(StopIteration)
+        await self.recv_queue.put(StopAsyncIteration)
 
         if self._stream_writer:
             logger.debug('wait close')
@@ -57,12 +57,14 @@ async def disconnect(self):
 
         logger.debug('Disconnected.')
 
-    async def __aiter__(self):
-        while True:
-            result = await self.recv_queue.get()
-            if result is StopIteration:
-                raise StopAsyncIteration
-            yield result
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        result = await self.recv_queue.get()
+        if result is StopAsyncIteration:
+            raise result
+        return result
 
     async def put_message(self, message):
         await self.send_queue.put(message)
@@ -81,6 +83,10 @@ class TCPConnection(Connection):
     FMT = '<I4s'
     FMT_SIZE = struct.calcsize(FMT)
 
+    @property
+    def local_address(self):
+        return self._stream_writer.get_extra_info('sockname')[0]
+
     async def connect(self, server_addr):
         host, port = server_addr
 
@@ -108,9 +114,6 @@ async def _writer_loop_func(self):
             try:
                 self._stream_writer.write(packet)
                 await self._stream_writer.drain()
-            except asyncio.CancelledError:
-                logger.debug('Writer loop was cancelled.')
-                break
             except Exception as e:
                 logger.debug(f'Writer error: {e}')
                 await self.disconnect()
@@ -128,9 +131,6 @@ async def _reader_loop_func(self):
                 logger.debug(f'Received data: {data}')
                 self._readbuf += data
                 await self._read_packets()
-            except asyncio.CancelledError:
-                logger.debug('Reader loop was cancelled.')
-                break
             except Exception as e:
                 logger.debug(f'Reader error: {e}')
                 await self.disconnect()
@@ -195,6 +195,13 @@ async def connect(self, server_addr: tuple[str, int]):
             logger.error(f'WebSocket connection failed: {e}')
             return False
 
+    async def disconnect(self):
+        if self.ws is not None:
+            logger.debug('Disconnecting WebSocket...')
+            await self.ws.close()
+
+        await super().disconnect()
+
     async def _writer_loop_func(self):
         while True:
             message = await self.send_queue.get()
@@ -221,10 +228,3 @@ async def _reader_loop_func(self):
                 logger.debug(f'Reader error: {e}')
                 await self.disconnect()
                 return
-
-    async def disconnect(self):
-        if self.ws is not None:
-            logger.debug('Disconnecting WebSocket...')
-            await self.ws.close()
-
-        await super().disconnect()