/
/
/
1"""VBAN subclasses to prevent unnecessary packet processing when the plugin isn't in use."""
2
3from __future__ import annotations
4
5import asyncio
6from dataclasses import dataclass
7from typing import TYPE_CHECKING, Any
8
9from aiovban.asyncio import AsyncVBANClient
10from aiovban.asyncio.protocol import VBANListenerProtocol
11
12if TYPE_CHECKING:
13 from .provider import VBANReceiverProvider
14
15
16@dataclass
17class VBANListenerProtocolMod(VBANListenerProtocol): # type: ignore[misc]
18 """VBANListenerProcotol workaround."""
19
20 controller: VBANReceiverProvider | None = None
21
22 def datagram_received(self, data: bytes, addr: tuple[str, int]) -> None:
23 """Handle received datagram."""
24 # No need to process the datagram if no queue is currently streaming us
25 if self.controller and not self.controller._in_use_by_player:
26 return
27 super().datagram_received(data, addr)
28
29
30@dataclass
31class AsyncVBANClientMod(AsyncVBANClient): # type: ignore[misc]
32 """AsyncVBANClient workaround."""
33
34 async def listen(
35 self,
36 address: str = "0.0.0.0",
37 port: int = 6980,
38 loop: asyncio.AbstractEventLoop | None = None,
39 controller: VBANReceiverProvider | None = None,
40 ) -> asyncio.Future[Any]:
41 """Create UDP listener."""
42 loop = loop or asyncio.get_running_loop()
43
44 # Create a socket and set the options
45 self._transport, proto = await loop.create_datagram_endpoint(
46 lambda: VBANListenerProtocolMod(self, controller),
47 local_addr=(address, port),
48 allow_broadcast=True,
49 )
50 return proto.done # type: ignore[no-any-return]
51