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()
|
