From 125afc89073b4fc69a3f42650b3d4f6fa6ccaa47 Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Fri, 11 Apr 2025 13:34:05 +0200
Subject: [PATCH] fix(rpc_server): enhance serialization logic for BECConnector
 objects and fix return types

---
 bec_widgets/utils/rpc_server.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/bec_widgets/utils/rpc_server.py b/bec_widgets/utils/rpc_server.py
index a3437bb4..1b831e05 100644
--- a/bec_widgets/utils/rpc_server.py
+++ b/bec_widgets/utils/rpc_server.py
@@ -4,7 +4,7 @@ import functools
 import traceback
 import types
 from contextlib import contextmanager
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 from bec_lib.client import BECClient
 from bec_lib.endpoints import MessageEndpoints
@@ -28,6 +28,9 @@ else:
 logger = bec_logger.logger
 
 
+T = TypeVar("T")
+
+
 @contextmanager
 def rpc_exception_hook(err_func):
     """This context replaces the popup message box for error display with a specific hook"""
@@ -143,13 +146,22 @@ class RPCServer:
                 res = self.serialize_object(res)
             return res
 
-    def serialize_object(self, obj):
-        if isinstance(obj, BECConnector):
-            # Respect RPC = False
-            if hasattr(obj, "RPC") and obj.RPC is False:
-                return None
-            return self._serialize_bec_connector(obj)
-        return obj
+    def serialize_object(self, obj: T) -> None | dict | T:
+        """
+        Serialize all BECConnector objects.
+
+        Args:
+            obj: The object to be serialized.
+
+        Returns:
+            None | dict | T: The serialized object or None if the object is not a BECConnector.
+        """
+        if not isinstance(obj, BECConnector):
+            return obj
+        # Respect RPC = False
+        if getattr(obj, "RPC", True) is False:
+            return None
+        return self._serialize_bec_connector(obj)
 
     def emit_heartbeat(self):
         logger.trace(f"Emitting heartbeat for {self.gui_id}")