apply cloud connection fixes

Author: TheCommCraft

scratchattach/cloud/_base.py

@@ -39,6 +39,7 @@ class EventStream(SupportsRead[Iterator[dict[str, Any]]], SupportsClose):
     """
     Allows you to stream events
     """
+    timeout: Optional[Union[float, int]] = None
 
 class AnyCloud(ABC, Generic[T]):
     """
@@ -129,38 +130,57 @@ def __init__(self, cloud: BaseCloud):
             warnings.warn("Initial cloud connection attempt failed, retrying...", exceptions.UnexpectedWebsocketEventWarning)
         self.packets_left = []
 
-    def receive_new(self, non_blocking: bool = False):
+    def receive_new(self, non_blocking: bool = False, timeout: Optional[float] = 0):
+        timeout = None if timeout is None else max(timeout, 0)
+        timeout_value = self.timeout if timeout is None else timeout
         if non_blocking:
-            self.source_cloud.websocket.settimeout(0)
-            try:
-                received = self.source_cloud.websocket.recv().splitlines()
-                self.packets_left.extend(received)
-            except Exception:
-                pass
+            timeout_value = 0
+        if self.source_cloud.websocket.gettimeout() != timeout_value:
+            self.source_cloud.websocket.settimeout(timeout_value)
+        # print("Receiving...")
+        try:
+            received = self.source_cloud.websocket.recv().splitlines()
+        except websocket.WebSocketTimeoutException:
             return
-        self.source_cloud.websocket.settimeout(None)
-        received = self.source_cloud.websocket.recv().splitlines()
+        # print(f"{received=}")
         self.packets_left.extend(received)
-    
+
     def read(self, amount: int = -1) -> Iterator[dict[str, Any]]:
+        # print("Reading...")
         i = 0
+        recv_once = amount == -1
+        recv_at_least = max(amount, 0)
+        start_time = time.time()
+        if self.timeout is not None:
+            has_timeout = True
+            timeout_end = start_time + self.timeout
+        else:
+            has_timeout = False
+            timeout_end = 0.0
+        done = False
+        # print("Getting data...")
         with self.reading:
-            try:
-                self.receive_new(amount != -1)
-                while (self.packets_left and amount == -1) or (amount != -1 and i < amount):
-                    if not self.packets_left and amount != -1:
-                        self.receive_new()
-                    yield json.loads(self.packets_left.pop(0))
-                    i += 1
-            except Exception:
-                self.source_cloud.reconnect()
-                self.receive_new(amount != -1)
-                while (self.packets_left and amount == -1) or (amount != -1 and i < amount):
-                    if not self.packets_left and amount != -1:
-                        self.receive_new()
-                    yield json.loads(self.packets_left.pop(0))
-                    i += 1
+            # print("Getting data...", end_time is None, end_time > time.time(), end_time is None or end_time > time.time())
+            while not done:
+                # print("Getting data...")
+                try:
+                    self.receive_new(not recv_once, timeout = timeout_end - time.time() if has_timeout else None)
+                    while ((not has_timeout or time.time() < timeout_end)
+                        and ((recv_once and self.packets_left) or (not recv_once and i < recv_at_least))):
+                        if not self.packets_left and not recv_once:
+                            self.receive_new(timeout = timeout_end - time.time() if has_timeout else None)
+                        if not self.packets_left:
+                            continue
+                        yield json.loads(self.packets_left.pop(0))
+                        i += 1
+                    done = True
+                except Exception:
+                    # traceback.print_exc()
+                    self.source_cloud.reconnect()
 
+    def __del__(self):
+        self.close()
+        
     def close(self) -> None:
         self.source_cloud.disconnect()
 
@@ -305,6 +325,10 @@ def _handshake(self):
         self._send_packet(packet)
 
     def connect(self):
+        if self.websocket:
+            self.websocket.close()
+        if self.event_stream:
+            self.event_stream = None
         self.websocket = websocket.WebSocket(sslopt={"cert_reqs": ssl.CERT_NONE})
         self.websocket.connect(
             self.cloud_host,
@@ -329,6 +353,8 @@ def disconnect(self):
             self.websocket.close()
         except Exception:
             pass
+        if self.event_stream:
+            self.event_stream = None
 
     def _assert_valid_value(self, value):
         if not (value in [True, False, float('inf'), -float('inf')]):
@@ -445,6 +471,9 @@ def create_event_stream(self):
             raise ValueError("Cloud already has an event stream.")
         self.event_stream = WebSocketEventStream(self)
         return self.event_stream
+    
+    def __del__(self):
+        self.disconnect()
 
 class LogCloudMeta(ABCMeta):
     def __instancecheck__(cls, instance) -> bool:

scratchattach/eventhandlers/_base.py

@@ -5,16 +5,21 @@
 from threading import Thread
 from collections.abc import Callable
 import traceback
+from typing import Optional
 from scratchattach.utils.requests import requests
 from scratchattach.utils import exceptions
 
 class BaseEventHandler(ABC):
     _events: defaultdict[str, list[Callable]]
     _threaded_events: defaultdict[str, list[Callable]]
+    running: bool
+    _thread: Optional[Thread]
+    _call_threads: list[Thread]
 
     def __init__(self):
         self._thread = None
         self.running = False
+        self._call_threads = []
         self._events = defaultdict(list)
         self._threaded_events = defaultdict(list)
 
@@ -35,14 +40,18 @@ def start(self, *, thread=True, ignore_exceptions=True):
             else:
                 self._thread = None
                 self._updater()
-    
+
     def call_event(self, event_name, args : list = []):
         try:
+            # print(f"Calling for {event_name}...")
             if event_name in self._threaded_events:
                 for func in self._threaded_events[event_name]:
-                    Thread(target=func, args=args).start()
+                    thread = Thread(target=func, args=args)
+                    self._call_threads.append(thread)
+                    thread.start()
             if event_name in self._events:
                 for func in self._events[event_name]:
+                    # print(f"Called {func}.")
                     func(*args)
         except Exception as e:
             if self.ignore_exceptions:
@@ -54,31 +63,44 @@ def call_event(self, event_name, args : list = []):
                 except Exception:
                     print(e)
             else:
-                raise(e)
+                raise e
 
     @abstractmethod
     def _updater(self):
         pass
 
-    def stop(self):
+    def __del__(self):
+        self.stop()
+
+    def stop(self, wait_call_threads: bool = True):
         """
         Permanently stops the event handler.
         """
+        # print("Stopping event handler...")
         self.running = False
-        if self._thread is not None:
+        thread = self._thread
+        if thread is not None:
+            thread.join()
             self._thread = None
+        if not wait_call_threads:
+            return
+        for thread in self._call_threads:
+            thread.join()
 
     def pause(self):
         """
         Pauses the event handler.
         """
         self.running = False
+        thread = self._thread
+        if thread is not None:
+            thread.join()
 
     def resume(self):
         """
         Resumes the event handler.
         """
-        if self.running is False:
+        if not self.running:
             self.start()
 
     def event(self, function=None, *, thread=False):

scratchattach/eventhandlers/cloud_events.py

@@ -29,11 +29,12 @@ def _updater(self):
 
         self.call_event("on_ready")
 
-        if self.running is False:
+        if not self.running:
             return
-        while True:
+        while self.running:
             try:
-                while True:
+                while self.running:
+                    self.source_stream.timeout = 1
                     for data in self.source_stream.read():
                         try:
                             _a = cloud_activity.CloudActivity(timestamp=time.time()*1000, _session=self._session, cloud=self.cloud)