Fixed handlers overwriting (#55), improved WS API

examples/websocket_args_example.py

@@ -1,34 +1,47 @@
 import asyncio
 import logging
 from getpass import getpass
-from python_aternos import Client, atwss
 
+from typing import Tuple, Dict, Any
+
+from python_aternos import Client, Streams
+
+# Request credentials
 user = input('Username: ')
 pswd = getpass('Password: ')
 
+# Debug logging
 logs = input('Show detailed logs? (y/n) ').strip().lower() == 'y'
 if logs:
     logging.basicConfig(level=logging.DEBUG)
 
+# Authentication
 aternos = Client.from_credentials(user, pswd)
 
-s = aternos.list_servers()[0]
+server = aternos.list_servers()[0]
-socket = s.wss()
+socket = server.wss()
 
 
-@socket.wssreceiver(atwss.Streams.console, 'Server 1')
+# Handler for console messages
-async def console(msg, args):
+@socket.wssreceiver(Streams.console, ('Server 1',))
+async def console(
+        msg: Dict[Any, Any],
+        args: Tuple[str]) -> None:
+
     print(args[0], 'received', msg)
 
 
-async def main():
+# Main function
-    s.start()
+async def main() -> None:
+    server.start()
     await socket.connect()
     await asyncio.create_task(loop())
 
 
-async def loop():
+# Keepalive
+async def loop() -> None:
     while True:
-        await asyncio.sleep(1)
+        await asyncio.Future()
+
 
 asyncio.run(main())

examples/websocket_status_example.py

@@ -0,0 +1,55 @@
+import asyncio
+import logging
+from getpass import getpass
+
+from typing import Tuple, Dict, Any
+
+from python_aternos import Client, Streams
+
+# Request credentials
+user = input('Username: ')
+pswd = getpass('Password: ')
+
+# Debug logging
+logs = input('Show detailed logs? (y/n) ').strip().lower() == 'y'
+if logs:
+    logging.basicConfig(level=logging.DEBUG)
+
+# Authentication
+aternos = Client.from_credentials(user, pswd)
+
+server = aternos.list_servers()[0]
+socket = server.wss()
+
+
+# Handler for server status
+@socket.wssreceiver(Streams.status, ('Server 1',))
+async def state(
+        msg: Dict[Any, Any],
+        args: Tuple[str]) -> None:
+
+    print(args[0], 'received', len(msg), 'symbols')
+
+    server._info = msg
+    print(
+        args[0],
+        server.subdomain,
+        'is',
+        server.status
+    )
+
+
+# Main function
+async def main() -> None:
+    server.start()
+    await socket.connect()
+    await asyncio.create_task(loop())
+
+
+# Keepalive
+async def loop() -> None:
+    while True:
+        await asyncio.Future()
+
+
+asyncio.run(main())

python_aternos/atwss.py

@@ -6,8 +6,9 @@ import json
 import asyncio
 import logging
 
+from typing import Iterable
 from typing import Union, Any
-from typing import Tuple, Dict
+from typing import Tuple, List, Dict
 from typing import Callable, Coroutine
 from typing import TYPE_CHECKING
 
@@ -66,12 +67,19 @@ class AternosWss:
         cookies = atserv.atconn.session.cookies
         self.session = cookies['ATERNOS_SESSION']
 
-        recvtype = Dict[Streams, ArgsTuple]
+        self.recv: Dict[Streams, List[ArgsTuple]]
-        self.recv: recvtype = {}
+        self.recv = {
+            Streams.status: [],
+            Streams.queue: [],
+            Streams.console: [],
+            Streams.ram: [],
+            Streams.tps: [],
+        }
+
         self.autoconfirm = autoconfirm
         self.confirmed = False
 
-        self.socket: Any
+        self.socket: Any = None
         self.keep: asyncio.Task
         self.msgs: asyncio.Task
 
@@ -86,7 +94,7 @@ class AternosWss:
     def wssreceiver(
             self,
             stream: Streams,
-            *args: Any) -> Callable[[FunctionT], Any]:
+            arg: Tuple[Any, ...] = ()) -> Callable[[FunctionT], Any]:
 
         """Decorator that marks your function as a stream receiver.
         When websocket receives message from the specified stream,
@@ -94,14 +102,26 @@ class AternosWss:
 
         Args:
             stream (Streams): Stream that your function should listen
-            *args (tuple, optional): Arguments which will be passed to your function
+            arg (Tuple[Any, ...], optional): Arguments which will be passed to your function
 
         Returns:
             ...
         """
 
-        def decorator(func: FunctionT) -> None:
+        def decorator(func: FunctionT) -> Callable[[Any, Any], Coroutine[Any, Any, Any]]:
-            self.recv[stream] = (func, args)
+
+            handlers = self.recv.get(stream, None)
+
+            if handlers is None:
+                self.recv[stream] = [(func, arg)]
+            else:
+                handlers.append((func, arg))
+
+            async def wrapper(*args, **kwargs) -> Any:
+                return await func(*args, **kwargs)
+
+            return wrapper
+
         return decorator
 
     async def connect(self) -> None:
@@ -173,6 +193,10 @@ class AternosWss:
                     if not isinstance(strm, Streams):
                         continue
 
+                    # If the handlers list is empty
+                    if not self.recv.get(strm):
+                        continue
+
                     if strm.stream:
                         logging.debug('Requesting %s stream', strm.stream)
                         await self.send({
@@ -200,10 +224,32 @@ class AternosWss:
                 Message, may be a string or a dict
         """
 
+        if self.socket is None:
+            logging.warning('Did you forget to call socket.connect?')
+            await self.connect()
+
         if isinstance(obj, dict):
             obj = json.dumps(obj)
 
         await self.socket.send(obj)
+    
+    async def command(self, cmd: str) -> None:
+
+        """Sends a Minecraft command
+        to the websocket server
+
+        Args:
+            cmd (str): Command, slash at
+                the beginning is not required
+        """
+
+        await self.send(
+            {
+                'stream': 'console',
+                'type': 'command',
+                'data': cmd,
+            }
+        )
 
     async def wssworker(self) -> None:
 
@@ -216,7 +262,8 @@ class AternosWss:
 
     async def keepalive(self) -> None:
 
-        """Each 49 seconds sends keepalive ping to websocket server"""
+        """Each 49 seconds sends keepalive ping
+        to the websocket server"""
 
         try:
             while True:
@@ -256,19 +303,24 @@ class AternosWss:
 
                 if msgtype in self.recv:
 
-                    # function info tuple
+                    # function info tuples
-                    func: ArgsTuple = self.recv[msgtype]
+                    handlers: Iterable[ArgsTuple]
+                    handlers = self.recv.get(msgtype, ())
 
-                    # if arguments is not empty
+                    for func in handlers:
-                    if func[1]:
+
-                        # call the function with args
+                        # if arguments is not empty
-                        coro = func[0](msg, func[1])  # type: ignore
+                        if func[1]:
-                    else:
+                            # call the function with args
-                        # mypy error: Too few arguments
+                            coro = func[0](msg, func[1])  # type: ignore
-                        # looks like bug, so it is ignored
+
-                        coro = func[0](msg)  # type: ignore
+                        else:
-                    # run
+                            # mypy error: too few arguments
-                    asyncio.create_task(coro)
+                            # looks like a bug, so it is ignored
+                            coro = func[0](msg)  # type: ignore
+
+                        # run
+                        asyncio.create_task(coro)
 
             except asyncio.CancelledError:
                 break