Merge pull request #21 from ThatXliner/v2

Version 2

.github/workflows/release.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: 3
       - name: Install Poetry
@@ -20,18 +20,43 @@ jobs:
         with:
           upload_url: ${{ github.event.release.upload_url }}
           asset_path: dist/*
+  release:
+    runs-on: ubuntu-latest
+    needs:
+      - build
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Sign the dists with Sigstore
+        uses: sigstore/[email protected]
+        with:
+          inputs: >-
+            ./dist/*.tar.gz
+            ./dist/*.whl
+      - name: Store the distribution packages
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
   upload:
     runs-on: ubuntu-latest
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/aioudp # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+    needs:
+      - build
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v4
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
         with:
-          python-version: 3
-      - name: Install Poetry
-        uses: snok/install-poetry@v1
-      - name: Build artifacts
-        run: poetry build
-      - name: Upload artifacts to PyPi
-        run: poetry publish --username __token__ --password $PYPI_TOKEN
-        env:
-          PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/

README.md

@@ -1,5 +1,16 @@
 # AioUDP
 
+<picture>
+  <!-- Dark mode image -->
+  <source srcset="docs/assets/[email protected]" media="(prefers-color-scheme: dark)">
+  <!-- Light mode image -->
+  <source srcset="docs/assets/[email protected]" media="(prefers-color-scheme: light)">
+  <!-- Fallback image -->
+  <img src="docs/assets/[email protected]" alt="AioUDP Banner" style="max-width: 100%; height: auto;">
+</picture>
+
+> A better API for asynchronous UDP
+
 [![Poetry](https://img.shields.io/endpoint?url=https://python-poetry.org/badge/v0.json)](https://python-poetry.org/)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
@@ -13,7 +24,6 @@
 [![PyPI](https://img.shields.io/pypi/v/aioudp)](https://pypi.org/project/aioudp)
 [![PyPI - License](https://img.shields.io/pypi/l/aioudp)](#license)
 
-> A better API for asynchronous UDP
 
 A [websockets](https://websockets.readthedocs.io/en/stable/index.html)-like API for [UDP](https://en.wikipedia.org/wiki/User_Datagram_Protocol)
 
@@ -31,10 +41,10 @@ async def main():
         async for message in connection:
             await connection.send(message)
 
-    # Optional. This is for properly exiting the server when Ctrl-C is pressed
-    # or when the process is killed/terminated
     loop = asyncio.get_running_loop()
     stop = loop.create_future()
+    # Optional. This is for properly exiting the server when Ctrl-C is pressed
+    # or when the process is killed/terminated
     loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
     loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
 

aioudp/__init__.py

@@ -8,4 +8,4 @@
 __all__ = ["connect", "serve", "Connection", "exceptions"]
 
 
-__version__ = "1.0.1"
+__version__ = "2.0.0"

aioudp/client.py

@@ -1,31 +1,43 @@
 """Client-side UDP connection."""
+
 from __future__ import annotations
 
 import asyncio
-import functools
 from contextlib import asynccontextmanager
-from dataclasses import dataclass
-from typing import AsyncIterator, NoReturn
+from dataclasses import dataclass, field
+from typing import AsyncIterator
 
 from aioudp import connection
 
 
 @dataclass
 class _ClientProtocol(asyncio.DatagramProtocol):
-    msg_queue: asyncio.Queue[None | bytes]
+    on_connection: asyncio.Future[connection.Connection]
+    on_connection_lost: asyncio.Future[bool]
+    msg_queue: asyncio.Queue[bytes] = field(default_factory=asyncio.Queue)
+
+    def connection_made(self, transport: asyncio.DatagramTransport) -> None:
+        self.on_connection.set_result(
+            connection.Connection(
+                send_func=transport.sendto,
+                recv_func=self.msg_queue.get,
+                is_closing=transport.is_closing,
+                get_local_addr=lambda: transport.get_extra_info("sockname"),
+                get_remote_addr=lambda: transport.get_extra_info("peername"),
+            ),
+        )
 
-    def datagram_received(self, data: bytes, _: connection.AddrType) -> None:
+    def datagram_received(self, data: bytes, _addr: connection.AddrType) -> None:
         self.msg_queue.put_nowait(data)
 
-    def error_received(self, exc: Exception) -> NoReturn:
-        # Haven't figured out why this can happen
+    def error_received(self, exc: BaseException) -> None:
         raise exc
 
-    def connection_lost(self, exc: None | Exception) -> None:
-        # Haven't figured out why this can happen
-        if exc is not None:
+    def connection_lost(self, exc: Exception | None) -> None:
+        self.msg_queue.shutdown(immediate=True)
+        self.on_connection_lost.set_result(True)
+        if exc:
             raise exc
-        self.msg_queue.put_nowait(None)
 
 
 @asynccontextmanager
@@ -47,25 +59,16 @@ async def connect(host: str, port: int) -> AsyncIterator[connection.Connection]:
 
     """
     loop = asyncio.get_running_loop()
-    msgs: asyncio.Queue[None | bytes] = asyncio.Queue()
-    transport: asyncio.DatagramTransport
-    _: _ClientProtocol
+    on_connection = loop.create_future()
+    on_connection_lost = loop.create_future()
     transport, _ = await loop.create_datagram_endpoint(
-        lambda: _ClientProtocol(msgs),
+        lambda: _ClientProtocol(on_connection, on_connection_lost),
         remote_addr=(host, port),
     )
-    conn = connection.Connection(  # TODO(ThatXliner): REFACTOR: minimal args
-        # https://github.com/ThatXliner/aioudp/issues/15
-        send_func=transport.sendto,
-        recv_func=msgs.get,
-        is_closing=transport.is_closing,
-        get_local_addr=functools.partial(transport.get_extra_info, "sockname"),
-        get_remote_addr=functools.partial(transport.get_extra_info, "peername"),
-    )
+
+    conn = await on_connection
     try:
-        # This is to make sure that the connection works
-        # See https://github.com/ThatXliner/aioudp/pull/3 for more information
-        await conn.send(b"trash")
         yield conn
     finally:
         transport.close()
+        await on_connection_lost

aioudp/connection.py

@@ -1,6 +1,8 @@
 """Connection class for aioudp."""
+
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 from typing import Awaitable, Callable, Tuple
 
@@ -17,7 +19,7 @@ class Connection:  # TODO(ThatXliner): REFACTOR: minimal args
     """Represents a server-client connection. Do not instantiate manually."""
 
     send_func: Callable[[bytes], None]
-    recv_func: Callable[[], Awaitable[None | bytes]]
+    recv_func: Callable[[], Awaitable[bytes]]
     is_closing: Callable[[], bool]
     get_local_addr: Callable[[], AddrType]
     get_remote_addr: Callable[[], None | AddrType]
@@ -32,7 +34,7 @@ def local_address(self) -> AddrType:
 
         Returns
         -------
-            AddrType: This is a `tuple` containing the hostname and port
+            tuple[str, int]: This is a `tuple` containing the hostname and port
 
         """
         return self.get_local_addr()
@@ -46,7 +48,7 @@ def remote_address(self) -> None | AddrType:
 
         Returns
         -------
-            AddrType: This is a `tuple` containing the hostname and port
+            tuple[str, int]: This is a `tuple` containing the hostname and port
 
         """
         return self.get_remote_addr()
@@ -63,12 +65,10 @@ async def recv(self) -> bytes:
             exceptions.ConnectionClosedError: The connection is closed
 
         """
-        the_next_one = await self.recv_func()
-        if the_next_one is None:
-            assert self.is_closing()
-            msg = "The connection is closed"
-            raise exceptions.ConnectionClosedError(msg)
-        return the_next_one
+        try:
+            return await self.recv_func()
+        except asyncio.QueueShutDown:
+            raise exceptions.ConnectionClosedError  # noqa: B904
 
     async def send(self, data: bytes) -> None:
         """Send a message to the connection.
@@ -77,7 +77,7 @@ async def send(self, data: bytes) -> None:
             Since this is UDP, there is no guarantee that the message will be sent
 
         Args:
-        ----
+        -----
             data (bytes): The message in bytes to send
 
         Raises:

aioudp/server.py

@@ -1,19 +1,20 @@
 """Server-side UDP connection."""
+
 from __future__ import annotations
 
 import asyncio
 import functools
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
-from typing import AsyncIterator, Awaitable, Callable, Coroutine, NoReturn
+from typing import Any, AsyncIterator, Callable, Coroutine, NoReturn
 
 from aioudp import connection
 
 
 @dataclass
 class _ServerProtocol(asyncio.DatagramProtocol):
-    handler: Callable[[connection.Connection], Coroutine[NoReturn, NoReturn, None]]
-    msg_queues: dict[connection.AddrType, asyncio.Queue[None | bytes]] = field(
+    handler: Callable[[connection.Connection], Coroutine[Any, Any, None]]
+    msg_queues: dict[connection.AddrType, asyncio.Queue[bytes]] = field(
         default_factory=dict,
     )
     msg_handlers: dict[
@@ -26,7 +27,7 @@ class _ServerProtocol(asyncio.DatagramProtocol):
 
     def connection_made(
         self,
-        transport: asyncio.transports.DatagramTransport,  # type: ignore[override]
+        transport: asyncio.DatagramTransport,  # type: ignore[override]
         # I am aware of the Liskov subsitution principle
         # but asyncio.DatagramProtocol had this function signature
     ) -> None:
@@ -53,31 +54,38 @@ def done(_) -> None:  # type: ignore[no-untyped-def]  # noqa: ANN001
                             self.transport.get_extra_info,
                             "sockname",
                         ),
-                        get_remote_addr=lambda: addr,
+                        # This should theoretically be
+                        # the same as the `addr` parameter
+                        get_remote_addr=functools.partial(
+                            self.transport.get_extra_info,
+                            "peername",
+                        ),
                     ),
                 ),
             )
             self.msg_handlers[addr].add_done_callback(done)
+
         self.msg_queues[addr].put_nowait(data)
 
     def error_received(self, exc: Exception) -> NoReturn:
         # Haven't figured out why this can happen
         raise exc
 
+    # The server is done
     def connection_lost(self, exc: None | Exception) -> None:
+        for key in self.msg_queues:
+            self.msg_queues[key].shutdown(immediate=True)
+            self.msg_handlers[key].cancel()
         # Haven't figured out why this can happen
         if exc is not None:
             raise exc
-        for key in self.msg_queues:
-            self.msg_queues[key].put_nowait(None)
-            self.msg_handlers[key].cancel()
 
 
 @asynccontextmanager
 async def serve(
     host: str,
     port: int,
-    handler: Callable[[connection.Connection], Awaitable[None]],
+    handler: Callable[[connection.Connection], Coroutine[Any, Any, None]],
 ) -> AsyncIterator[None]:
     """Run a UDP server.
 
@@ -99,16 +107,11 @@ async def serve(
             and doesn't need to return anything.
 
     """
-
-    async def wrap_handler(con: connection.Connection) -> None:
-        await con.recv()
-        return await handler(con)
-
     loop = asyncio.get_running_loop()
     transport: asyncio.BaseTransport
     _: asyncio.BaseProtocol
     transport, _ = await loop.create_datagram_endpoint(
-        lambda: _ServerProtocol(wrap_handler),
+        lambda: _ServerProtocol(handler),
         local_addr=(host, port),
     )
     try: