Merge pull request #48 from olijeffers0n/feat/async_ws

Feat- Async WebSocket

requirements.txt

@@ -1,4 +1,4 @@
-websocket_client
+websockets
 Pillow
 asyncio
 rustPlusPushReceiver==0.4.1

rustplus/__init__.py

@@ -22,5 +22,5 @@
 
 __name__ = "rustplus"
 __author__ = "olijeffers0n"
-__version__ = "5.5.13"
+__version__ = "5.6.0"
 __support__ = "Discord: https://discord.gg/nQqJe8qvP8"

rustplus/api/base_rust_api.py

@@ -39,7 +39,6 @@ def __init__(
         event_loop: asyncio.AbstractEventLoop = None,
         rate_limiter: RateLimiter = None,
     ) -> None:
-
         if ip is None:
             raise ValueError("Ip cannot be None")
         if steam_id is None:
@@ -80,16 +79,15 @@ async def _handle_ratelimit(self, amount=1) -> None:
         :return: None
         """
         while True:
-
-            if self.remote.ratelimiter.can_consume(self.server_id, amount):
-                self.remote.ratelimiter.consume(self.server_id, amount)
+            if await self.remote.ratelimiter.can_consume(self.server_id, amount):
+                await self.remote.ratelimiter.consume(self.server_id, amount)
                 break
 
             if self.raise_ratelimit_exception:
                 raise RateLimitError("Out of tokens")
 
             await asyncio.sleep(
-                self.remote.ratelimiter.get_estimated_delay_time(self.server_id, amount)
+                await self.remote.ratelimiter.get_estimated_delay_time(self.server_id, amount)
             )
 
         self.heartbeat.reset_rhythm()
@@ -151,7 +149,7 @@ async def close_connection(self) -> None:
 
         :return: None
         """
-        self.remote.close()
+        await self.remote.close()
 
     async def disconnect(self) -> None:
         """
@@ -173,7 +171,7 @@ async def send_wakeup_request(self) -> None:
         app_request.get_time = AppEmpty()
         app_request.get_time._serialized_on_wire = True
 
-        self.remote.ignored_responses.append(app_request.seq)
+        await self.remote.add_ignored_response(app_request.seq)
 
         await self.remote.send_message(app_request)
 
@@ -238,7 +236,8 @@ async def switch_server(
         self.remote.server_id = ServerID(ip, port, steam_id, player_token)
 
         # reset ratelimiter
-        self.remote.ratelimiter.remove(self.server_id)
+        self.remote.use_proxy = use_proxy
+        await self.remote.ratelimiter.remove(self.server_id)
         self.remote.ratelimiter.add_socket(
             self.server_id,
             self.ratelimit_limit,
@@ -284,7 +283,6 @@ def command(
             return RegisteredListener(coro.__name__, cmd_data.coro)
 
         def wrap_func(coro):
-
             if self.command_options is None:
                 raise CommandsNotEnabledError("Not enabled")
 
@@ -341,7 +339,6 @@ def entity_event(self, eid):
         """
 
         def wrap_func(coro) -> RegisteredListener:
-
             if isinstance(coro, RegisteredListener):
                 coro = coro.get_coro()
 
@@ -588,14 +585,14 @@ async def get_tc_storage_contents(
         """
         raise NotImplementedError("Not Implemented")
 
-    async def get_camera_manager(self, id: str) -> CameraManager:
+    async def get_camera_manager(self, cam_id: str) -> CameraManager:
         """
         Gets a camera manager for a given camera ID
 
         NOTE: This will override the current camera manager if one exists for the given ID so you cannot have multiple
 
-        :param id: The ID of the camera
+        :param cam_id: The ID of the camera
         :return CameraManager: The camera manager
-        :raises RequestError: If the camera is not found or you cannot access it. See reason for more info
+        :raises RequestError: If the camera is not found, or you cannot access it. See reason for more info
         """
         raise NotImplementedError("Not Implemented")

rustplus/api/remote/camera/camera_manager.py

@@ -1,10 +1,9 @@
 import time
-from typing import Iterable, Union, List, Coroutine, TypeVar, Set
+from typing import Iterable, Union, List, Coroutine, TypeVar, Set, Callable
 
 from PIL import Image
 
 from .camera_parser import Parser
-from ..events import EventLoopManager, EventHandler
 from ..rustplus_proto import (
     AppCameraInput,
     Vector2,
@@ -31,31 +30,29 @@ def __init__(
             self._cam_info_message.width, self._cam_info_message.height
         )
         self.time_since_last_subscribe: float = time.time()
-        self.frame_callbacks: Set[Coroutine] = set()
+        self.frame_callbacks: Set[Callable[[Image.Image], Coroutine]] = set()
 
-    def add_packet(self, packet) -> None:
+    async def add_packet(self, packet) -> None:
         self._last_packets.add(packet)
 
         if len(self.frame_callbacks) == 0:
             return
 
-        frame = self._create_frame()
+        frame = await self._create_frame()
 
         for callback in self.frame_callbacks:
-            EventHandler.schedule_event(
-                EventLoopManager.get_loop(self.rust_socket.server_id),
-                callback,
-                frame,
-            )
+            await callback(frame)
 
-    def on_frame_received(self, coro: Coroutine) -> Coroutine:
+    def on_frame_received(
+        self, coro: Callable[[Image.Image], Coroutine]
+    ) -> Callable[[Image.Image], Coroutine]:
         self.frame_callbacks.add(coro)
         return coro
 
     def has_frame_data(self) -> bool:
         return self._last_packets is not None and len(self._last_packets) > 0
 
-    def _create_frame(
+    async def _create_frame(
         self,
         render_entities: bool = True,
         entity_render_distance: float = float("inf"),
@@ -96,7 +93,7 @@ async def get_frame(
         entity_render_distance: float = float("inf"),
         max_entity_amount: int = float("inf"),
     ) -> Union[Image.Image, None]:
-        return self._create_frame(
+        return await self._create_frame(
             render_entities, entity_render_distance, max_entity_amount
         )
 
@@ -115,7 +112,6 @@ async def send_mouse_movement(self, mouse_delta: Vector) -> None:
     async def send_combined_movement(
         self, movements: Iterable[int] = None, joystick_vector: Vector = None
     ) -> None:
-
         if joystick_vector is None:
             joystick_vector = Vector()
 
@@ -138,7 +134,7 @@ async def send_combined_movement(
         app_request.camera_input = cam_input
 
         await self.rust_socket.remote.send_message(app_request)
-        self.rust_socket.remote.ignored_responses.append(app_request.seq)
+        await self.rust_socket.remote.add_ignored_response(app_request.seq)
 
     async def exit_camera(self) -> None:
         await self.rust_socket._handle_ratelimit()
@@ -147,7 +143,7 @@ async def exit_camera(self) -> None:
         app_request.camera_unsubscribe._serialized_on_wire = True
 
         await self.rust_socket.remote.send_message(app_request)
-        self.rust_socket.remote.ignored_responses.append(app_request.seq)
+        await self.rust_socket.remote.add_ignored_response(app_request.seq)
 
         self._open = False
         self._last_packets.clear()

rustplus/api/remote/camera/camera_parser.py

@@ -54,7 +54,6 @@ def reset_output(self) -> None:
         )
 
     def handle_camera_ray_data(self, data) -> None:
-
         if data is None:
             return
 
@@ -66,7 +65,6 @@ def handle_camera_ray_data(self, data) -> None:
         self._ray_lookback = [[0 for _ in range(3)] for _ in range(64)]
 
     def step(self) -> None:
-
         if self._rays is None:
             return
 
@@ -80,7 +78,6 @@ def process_rays_batch(self) -> bool:
             return True
 
         for h in range(100):
-
             if self.data_pointer >= len(self._rays.ray_data) - 1:
                 return True
 
@@ -107,7 +104,6 @@ def process_rays_batch(self) -> bool:
                 not (distance == 1 and alignment == 0 and material == 0)
                 and material != 7
             ):
-
                 self.colour_output[
                     x : x + self.scale_factor, y : y + self.scale_factor
                 ] = MathUtils._convert_colour(
@@ -153,7 +149,6 @@ def next_ray(self, ray_data) -> List[Union[float, int]]:
             self._ray_lookback[u][2] = i
 
         else:
-
             c = 192 & byte
 
             if c == 0:
@@ -217,7 +212,6 @@ def handle_entities(
         entity_render_distance: float,
         max_entity_amount: int,
     ) -> Any:
-
         image_data = np.array(image_data)
 
         players = [player for player in entities if player.type == 2]
@@ -265,7 +259,6 @@ def handle_entities(
         text = set()
 
         for entity in entities:
-
             if entity.position.z > entity_render_distance and entity.type == 1:
                 continue
 
@@ -307,7 +300,6 @@ def handle_entity(
         aspect_ratio,
         text,
     ) -> None:
-
         entity.size.x = min(entity.size.x, 5)
         entity.size.y = min(entity.size.y, 5)
         entity.size.z = min(entity.size.z, 5)
@@ -418,7 +410,6 @@ def render(
         entity_render_distance: float,
         max_entity_amount: int,
     ) -> Image.Image:
-
         # We have the output array filled with RayData objects
         # We can get the material at each pixel and use that to get the colour
         # We can then use the alignment to get the alpha value
@@ -556,7 +547,6 @@ def _convert_colour(
         cls,
         colour: Tuple[float, float, float, float],
     ) -> Tuple[int, int, int]:
-
         if colour in cls.COLOUR_CACHE:
             return cls.COLOUR_CACHE[colour]
 
@@ -589,7 +579,6 @@ def solve_quadratic(a: float, b: float, c: float, larger: bool) -> float:
 
     @classmethod
     def get_tree_vertices(cls, size) -> np.ndarray:
-
         if size in cls.VERTEX_CACHE:
             return cls.VERTEX_CACHE[size]
 
@@ -599,7 +588,6 @@ def get_tree_vertices(cls, size) -> np.ndarray:
         vertex_list = []
 
         for x_value in [size.y / 8, -size.y / 8]:
-
             for i in range(number_of_segments):
                 angle = segment_angle * i
 
@@ -616,7 +604,6 @@ def get_tree_vertices(cls, size) -> np.ndarray:
 
     @classmethod
     def get_player_vertices(cls, size) -> np.ndarray:
-
         if size in cls.VERTEX_CACHE:
             return cls.VERTEX_CACHE[size]
 
@@ -633,9 +620,7 @@ def get_player_vertices(cls, size) -> np.ndarray:
 
         x = 0
         while x <= width:
-
             for offset in range(-1, 2, 2):
-
                 x_value = x * offset
 
                 # Use the quadratic formula to find the y values of the top and bottom of the pill

rustplus/api/remote/events/event_handler.py

@@ -1,31 +1,16 @@
-import asyncio
-import logging
-from asyncio.futures import Future
-from typing import Set, Coroutine, Any
+from typing import Set, Union
 
 from ....utils import ServerID
 from .events import EntityEvent, TeamEvent, ChatEvent, ProtobufEvent
 from .registered_listener import RegisteredListener
-from .event_loop_manager import EventLoopManager
 from ..rustplus_proto import AppMessage
 
 
 class EventHandler:
     @staticmethod
-    def schedule_event(
-        loop: asyncio.AbstractEventLoop, coro: Coroutine, arg: Any
+    async def run_entity_event(
+        name: Union[str, int], app_message: AppMessage, server_id: ServerID
     ) -> None:
-        def callback(inner_future: Future):
-            if inner_future.exception() is not None:
-                logging.getLogger("rustplus.py").exception(inner_future.exception())
-
-        future: Future = asyncio.run_coroutine_threadsafe(coro(arg), loop)
-        future.add_done_callback(callback)
-
-    def run_entity_event(
-        self, name: str, app_message: AppMessage, server_id: ServerID
-    ) -> None:
-
         handlers: Set[RegisteredListener] = EntityEvent.handlers.get_handlers(
             server_id
         ).get(str(name))
@@ -36,40 +21,30 @@ def run_entity_event(
         for handler in handlers.copy():
             coro, event_type = handler.data
 
-            self.schedule_event(
-                EventLoopManager.get_loop(server_id),
-                coro,
-                EntityEvent(app_message, event_type),
-            )
-
-    def run_team_event(self, app_message: AppMessage, server_id: ServerID) -> None:
+            await coro(EntityEvent(app_message, event_type))
 
+    @staticmethod
+    async def run_team_event(app_message: AppMessage, server_id: ServerID) -> None:
         handlers: Set[RegisteredListener] = TeamEvent.handlers.get_handlers(server_id)
         for handler in handlers.copy():
             coro = handler.data
 
-            self.schedule_event(
-                EventLoopManager.get_loop(server_id), coro, TeamEvent(app_message)
-            )
-
-    def run_chat_event(self, app_message: AppMessage, server_id: ServerID) -> None:
+            await coro(TeamEvent(app_message))
 
+    @staticmethod
+    async def run_chat_event(app_message: AppMessage, server_id: ServerID) -> None:
         handlers: Set[RegisteredListener] = ChatEvent.handlers.get_handlers(server_id)
         for handler in handlers.copy():
             coro = handler.data
 
-            self.schedule_event(
-                EventLoopManager.get_loop(server_id), coro, ChatEvent(app_message)
-            )
-
-    def run_proto_event(self, byte_data: bytes, server_id: ServerID) -> None:
+            await coro(ChatEvent(app_message))
 
+    @staticmethod
+    async def run_proto_event(byte_data: bytes, server_id: ServerID) -> None:
         handlers: Set[RegisteredListener] = ProtobufEvent.handlers.get_handlers(
             server_id
         )
         for handler in handlers.copy():
             coro = handler.data
 
-            self.schedule_event(
-                EventLoopManager.get_loop(server_id), coro, ProtobufEvent(byte_data)
-            )
+            await coro(ProtobufEvent(byte_data))

rustplus/api/remote/events/event_loop_manager.py

@@ -5,7 +5,6 @@
 
 
 class EventLoopManager:
-
     _loop: Dict[ServerID, asyncio.AbstractEventLoop] = {}
 
     @staticmethod