From e709afd657eda10bd1f0373419c1be0f66167cb8 Mon Sep 17 00:00:00 2001
From: HeadTriXz <l.knol@outlook.com>
Date: Wed, 29 May 2024 23:00:32 +0200
Subject: [PATCH] feat: check for active connections

---
 src/lane_assist/lane_assist.py                |  2 +-
 src/lane_assist/preprocessing/generator.py    |  2 +-
 src/simulation/main.py                        |  2 +-
 src/telemetry/app.py                          | 23 ++++++++++++-------
 .../data_stream/websocket_handler.py          |  8 ++++++-
 5 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/src/lane_assist/lane_assist.py b/src/lane_assist/lane_assist.py
index b86f68a..acaabca 100644
--- a/src/lane_assist/lane_assist.py
+++ b/src/lane_assist/lane_assist.py
@@ -116,7 +116,7 @@ def lane_assist_loop(self, image: np.ndarray) -> None:
         self.__stop_line_assist.detect_and_handle(image, filtered_lines)
 
         # If telemetry is enabled, send the image to the telemetry server.
-        if config["telemetry"]["enabled"]:
+        if config["telemetry"]["enabled"] and self.telemetry.any_listening():
             rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
 
             # Draw the path on the image.
diff --git a/src/lane_assist/preprocessing/generator.py b/src/lane_assist/preprocessing/generator.py
index f821e7f..26d7519 100644
--- a/src/lane_assist/preprocessing/generator.py
+++ b/src/lane_assist/preprocessing/generator.py
@@ -51,7 +51,7 @@ def __generator() -> Generator[np.ndarray, None, None]:
             topdown = calibration.transform([left_image, center_image, right_image])
             thresholded = cv2.threshold(topdown, config["preprocessing"]["white_threshold"], 255, cv2.THRESH_BINARY)[1]
 
-            if config["telemetry"]["enabled"]:
+            if config["telemetry"]["enabled"] and telemetry.any_listening():
                 telemetry.websocket_handler.send_image("left", left_image)
                 telemetry.websocket_handler.send_image("center", center_image)
                 telemetry.websocket_handler.send_image("right", right_image)
diff --git a/src/simulation/main.py b/src/simulation/main.py
index 235dd4d..77469c6 100644
--- a/src/simulation/main.py
+++ b/src/simulation/main.py
@@ -65,7 +65,7 @@ def get_sim_image_generator() -> Generator[np.ndarray, None, None]:
                 front_view, calibration.topdown_matrix, calibration.output_shape, flags=cv2.INTER_NEAREST
             )
 
-            if config["telemetry"]["enabled"]:
+            if config["telemetry"]["enabled"] and telemetry.any_listening():
                 telemetry.websocket_handler.send_image("topdown", top_view)
 
             yield top_view
diff --git a/src/telemetry/app.py b/src/telemetry/app.py
index 0205c12..50f935c 100644
--- a/src/telemetry/app.py
+++ b/src/telemetry/app.py
@@ -50,6 +50,21 @@ def __init__(self) -> None:
         self.__app.mount("/js", StaticFiles(directory=get_path("static/js")), name="static")
         self.__app.mount("/css", StaticFiles(directory=get_path("static/css")), name="static")
 
+    def add_callback_function(self, name: str, func: callable) -> None:
+        """Add a callback function.
+
+        :param name: The name of the callback function.
+        :param func: The callback function to add.
+        """
+        self.available_functions[name] = func
+
+    def any_listening(self) -> bool:
+        """Check if there are any active connections.
+
+        :return: Whether there are any active connections.
+        """
+        return self.websocket_handler.any_active()
+
     def start(self) -> None:
         """Start the telemetry server."""
         if config["telemetry"]["enabled"]:
@@ -68,14 +83,6 @@ def __index_route(self) -> HTMLResponse:
             html = f.read().replace("$root-url", f"{get_ip()}:{self.__port}")
             return HTMLResponse(content=html)
 
-    def add_callback_function(self, name: str, func: callable) -> None:
-        """Add a callback function.
-
-        :param name: The name of the callback function.
-        :param func: The callback function to add.
-        """
-        self.available_functions[name] = func
-
     def __execute_function(self, name: str) -> Any:
         """Execute a function based on the provided name.
 
diff --git a/src/telemetry/data_stream/websocket_handler.py b/src/telemetry/data_stream/websocket_handler.py
index bd08e3b..a8b5212 100644
--- a/src/telemetry/data_stream/websocket_handler.py
+++ b/src/telemetry/data_stream/websocket_handler.py
@@ -61,7 +61,6 @@ class WebsocketHandler:
     def __init__(self) -> None:
         """Initialize the websocket handler."""
         self.websocket_clients = {}
-        self.websockets_active = {}
 
     def add_socket(self, name: str, websocket: WebSocket, loop: AbstractEventLoop) -> WebsocketDataStream:
         """Add a websocket client to the list of clients.
@@ -77,6 +76,13 @@ def add_socket(self, name: str, websocket: WebSocket, loop: AbstractEventLoop) -
         self.websocket_clients[name].append(WebsocketDataStream(websocket, loop))
         return self.websocket_clients[name][-1]
 
+    def any_active(self) -> bool:
+        """Check if there are any active websockets.
+
+        :return: Whether there are any active websockets.
+        """
+        return any(len(clients) > 0 for clients in self.websocket_clients.values())
+
     def send_image(self, name: str, image: np.ndarray) -> None:
         """Send image on channel with the given name.