Repository: eerimoq/asyncudp (language: python)

Issue:
Add async context support
Calling `close` can be forgotten. It's nice practise to have context manager to close resources.


Discussion / hints:
Feel free to implement one and submit a PR.

Fix patch:
diff --git a/asyncudp/__init__.py b/asyncudp/__init__.py
index 0dd21f8..5cba419 100644
--- a/asyncudp/__init__.py
+++ b/asyncudp/__init__.py
@@ -57,6 +57,12 @@ async def recvfrom(self):
 
         return await self._protocol.recvfrom()
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc_info):
+        self.close()
+
 
 async def create_socket(local_addr=None, remote_addr=None):
     """Create a UDP socket with given local and remote addresses.


Test patch:
diff --git a/tests/test_asyncudp.py b/tests/test_asyncudp.py
index 844355f..c504f4b 100644
--- a/tests/test_asyncudp.py
+++ b/tests/test_asyncudp.py
@@ -69,3 +69,32 @@ async def cancel(self):
         await asyncio.sleep(1.0)
         task.cancel()
         await event.wait()
+
+    def test_context(self):
+        asyncio.run(self.context())
+
+    async def context(self):
+        server_addr = ('127.0.0.1', 13000)
+        client_addr = ('127.0.0.1', 13001)
+
+        server = await asyncudp.create_socket(local_addr=server_addr)
+        client = await asyncudp.create_socket(local_addr=client_addr)
+
+        async with server, client:
+            client.sendto(b'local_addresses to server', server_addr)
+            data, addr = await server.recvfrom()
+
+            self.assertEqual(data, b'local_addresses to server')
+            self.assertEqual(addr, client_addr)
+
+            server.sendto(b'local_addresses to client', client_addr)
+            data, addr = await client.recvfrom()
+
+            self.assertEqual(data, b'local_addresses to client')
+            self.assertEqual(addr, server_addr)
+
+            self.assertEqual(server._transport.is_closing(), False)
+            self.assertEqual(client._transport.is_closing(), False)
+
+        self.assertEqual(server._transport.is_closing(), True)
+        self.assertEqual(client._transport.is_closing(), True)