1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
|
import asyncio
import contextlib
from typing import AsyncGenerator, cast
class EchoServerProtocol(asyncio.DatagramProtocol):
def connection_made(self, transport: asyncio.BaseTransport) -> None:
self.transport = cast(asyncio.DatagramTransport, transport)
def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
self.transport.sendto(data, addr)
class EchoServer:
async def close(self) -> None:
self.udp_server.transport.close()
async def listen(self, host: str = "127.0.0.1", port: int = 0) -> None:
loop = asyncio.get_event_loop()
# listen for UDP
transport, self.udp_server = await loop.create_datagram_endpoint(
EchoServerProtocol, local_addr=(host, port)
)
self.udp_address = transport.get_extra_info("sockname")
@contextlib.asynccontextmanager
async def run_echo_server() -> AsyncGenerator[EchoServer, None]:
server = EchoServer()
await server.listen()
try:
yield server
finally:
await server.close()
|