refactor: replace connect with connect_slot

bec_widgets/bec_dispatcher.py

@@ -2,11 +2,10 @@ import argparse
 import itertools
 import os
 from dataclasses import dataclass
-from threading import RLock
 from typing import Callable
 
 from bec_lib import BECClient
-from bec_lib.core import BECMessage, MessageEndpoints, ServiceConfig
+from bec_lib.core import BECMessage, ServiceConfig
 from bec_lib.core.redis_connector import RedisConsumerThreaded
 from PyQt5.QtCore import QObject, pyqtSignal
 
@@ -31,9 +30,6 @@ class _Connection:
 
 
 class _BECDispatcher(QObject):
-    new_scan = pyqtSignal(dict, dict)
-    scan_segment = pyqtSignal(dict, dict)
-
     def __init__(self, bec_config=None):
         super().__init__()
         self.client = BECClient()
@@ -44,35 +40,8 @@ class _BECDispatcher(QObject):
             bec_config = "bec_config.yaml"
 
         self.client.initialize(config=ServiceConfig(config_path=bec_config))
-
-        self._slot_signal_map = {"on_scan_segment": self.scan_segment, "on_new_scan": self.new_scan}
         self._connections = {}
 
-        self._scan_id = None
-        scan_lock = RLock()
-
-        def _scan_segment_cb(msg):
-            msg = BECMessage.ScanMessage.loads(msg.value)[0]
-            with scan_lock:
-                # TODO: use ScanStatusMessage instead?
-                scan_id = msg.content["scanID"]
-                if self._scan_id != scan_id:
-                    self._scan_id = scan_id
-                    self.new_scan.emit(msg.content, msg.metadata)
-            self.scan_segment.emit(msg.content, msg.metadata)
-
-        scan_segment_topic = MessageEndpoints.scan_segment()
-        self._scan_segment_thread = self.client.connector.consumer(
-            topics=scan_segment_topic, cb=_scan_segment_cb
-        )
-        self._scan_segment_thread.start()
-
-    def connect(self, widget):
-        for slot_name, signal in self._slot_signal_map.items():
-            slot = getattr(widget, slot_name, None)
-            if callable(slot):
-                signal.connect(slot)
-
     def connect_slot(self, slot: Callable, topic: str) -> None:
         """Connect widget's pyqt slot, so that it is called on new pub/sub topic message
 

bec_widgets/line_plot_legacy.py

@@ -5,10 +5,10 @@ from typing import Any
 import numpy as np
 import pyqtgraph
 import pyqtgraph as pg
-from PyQt5.QtCore import pyqtSlot
-from PyQt5.QtWidgets import QTableWidgetItem, QCheckBox
-
 from bec_lib import BECClient
+from bec_lib.core import MessageEndpoints
+from PyQt5.QtCore import pyqtSlot
+from PyQt5.QtWidgets import QCheckBox, QTableWidgetItem
 from pyqtgraph import mkBrush, mkColor, mkPen
 from pyqtgraph.Qt import QtCore, QtWidgets, uic
 from pyqtgraph.Qt.QtCore import pyqtSignal
@@ -54,10 +54,7 @@ class BasicPlot(QtWidgets.QWidget):
         self.plotter_scan_id = None
 
         # TODO to be moved to utils function
-        plotstyles = {
-            "symbol": "o",
-            "symbolSize": 10,
-        }
+        plotstyles = {"symbol": "o", "symbolSize": 10}
         color_list = ["#384c6b", "#e28a2b", "#5E3023", "#e41a1c", "#984e83", "#4daf4a"]
         color_list = BasicPlot.golden_angle_color(colormap="CET-R2", num=len(self.y_value_list))
 
@@ -336,9 +333,9 @@ class BasicPlot(QtWidgets.QWidget):
 
 if __name__ == "__main__":
     import argparse
-    from bec_widgets.bec_dispatcher import bec_dispatcher
 
     from bec_widgets import ctrl_c
+    from bec_widgets.bec_dispatcher import bec_dispatcher
 
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -354,7 +351,7 @@ if __name__ == "__main__":
     app = QtWidgets.QApplication([])
     ctrl_c.setup(app)
     plot = BasicPlot(y_value_list=value.signals)
-    bec_dispatcher.connect(plot)
+    bec_dispatcher.connect_slot(plot.on_scan_segment, MessageEndpoints.scan_segment())
     plot.show()
     # client.callbacks.register("scan_segment", plot, sync=False)
     app.exec_()

bec_widgets/scan2d_plot.py

@@ -1,5 +1,8 @@
+from threading import RLock
+
 import numpy as np
 import pyqtgraph as pg
+from bec_lib.core import MessageEndpoints
 from bec_lib.core.logger import bec_logger
 from PyQt5.QtCore import pyqtProperty, pyqtSlot
 
@@ -14,7 +17,10 @@ pg.setConfigOptions(background="w", foreground="k", antialias=True)
 class BECScanPlot2D(pg.GraphicsView):
     def __init__(self, parent=None, background="default"):
         super().__init__(parent, background)
-        bec_dispatcher.connect(self)
+        bec_dispatcher.connect_slot(self.on_scan_segment, MessageEndpoints.scan_segment())
+
+        self._scanID = None
+        self._scanID_lock = RLock()
 
         self._x_channel = ""
         self._y_channel = ""
@@ -33,8 +39,7 @@ class BECScanPlot2D(pg.GraphicsView):
         self.imageItem = pg.ImageItem()
         self.plot_item.addItem(self.imageItem)
 
-    @pyqtSlot(dict, dict)
-    def on_new_scan(self, _scan_segment, metadata):
+    def reset_plots(self, _scan_segment, metadata):
         # TODO: Do we reset in case of a scan type change?
         self.imageItem.clear()
 
@@ -79,6 +84,13 @@ class BECScanPlot2D(pg.GraphicsView):
 
     @pyqtSlot(dict, dict)
     def on_scan_segment(self, scan_segment, metadata):
+        # reset plots on scanID change
+        with self._scanID_lock:
+            scan_id = scan_segment["scanID"]
+            if self._scanID != scan_id:
+                self._scanID = scan_id
+                self.reset_plots(scan_segment, metadata)
+
         if not self.z_channel or metadata["scan_name"] != "grid_scan":
             return
 

bec_widgets/scan_plot.py

@@ -1,4 +1,5 @@
 import itertools
+from threading import RLock
 
 import pyqtgraph as pg
 from bec_lib.core import MessageEndpoints
@@ -17,24 +18,33 @@ COLORS = ["#fd7f6f", "#7eb0d5", "#b2e061", "#bd7ebe", "#ffb55a"]
 class BECScanPlot(pg.GraphicsView):
     def __init__(self, parent=None, background="default"):
         super().__init__(parent, background)
-        bec_dispatcher.connect(self)
+        bec_dispatcher.connect_slot(self.on_scan_segment, MessageEndpoints.scan_segment())
 
         self.view = pg.PlotItem()
         self.setCentralItem(self.view)
 
+        self._scanID = None
+        self._scanID_lock = RLock()
+
         self._x_channel = ""
         self._y_channel_list = []
 
         self.scan_curves = {}
         self.dap_curves = {}
 
-    @pyqtSlot(dict, dict)
-    def on_new_scan(self, _scan_segment, _metadata):
+    def reset_plots(self, _scan_segment, _metadata):
         for plot_curve in {**self.scan_curves, **self.dap_curves}.values():
             plot_curve.setData(x=[], y=[])
 
     @pyqtSlot(dict, dict)
-    def on_scan_segment(self, scan_segment, _metadata):
+    def on_scan_segment(self, scan_segment, metadata):
+        # reset plots on scanID change
+        with self._scanID_lock:
+            scan_id = scan_segment["scanID"]
+            if self._scanID != scan_id:
+                self._scanID = scan_id
+                self.reset_plots(scan_segment, metadata)
+
         if not self.x_channel:
             return
 

tests/test_scan_plot.py

@@ -58,7 +58,7 @@ def test_scan_plot_clears_data(qtbot):
         },
         {"scanID": "test", "scan_number": 1, "scan_report_devices": ["x"]},
     )
-    plot.on_new_scan({}, {})
+    plot.reset_plots({}, {})
     plot.on_scan_segment(
         {
             "data": {