From 7335aa9597a519371c3f02f21ea89edb06333eaf Mon Sep 17 00:00:00 2001 From: Ivan Usov Date: Thu, 12 Oct 2023 15:13:57 +0200 Subject: [PATCH] refactor: replace connect with connect_slot --- bec_widgets/bec_dispatcher.py | 33 +-------------------------------- bec_widgets/line_plot_legacy.py | 15 ++++++--------- bec_widgets/scan2d_plot.py | 18 +++++++++++++++--- bec_widgets/scan_plot.py | 18 ++++++++++++++---- tests/test_scan_plot.py | 2 +- 5 files changed, 37 insertions(+), 49 deletions(-) diff --git a/bec_widgets/bec_dispatcher.py b/bec_widgets/bec_dispatcher.py index 96c6def2..9a6a5a42 100644 --- a/bec_widgets/bec_dispatcher.py +++ b/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 diff --git a/bec_widgets/line_plot_legacy.py b/bec_widgets/line_plot_legacy.py index 6a96b872..3dd6d39d 100644 --- a/bec_widgets/line_plot_legacy.py +++ b/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_() diff --git a/bec_widgets/scan2d_plot.py b/bec_widgets/scan2d_plot.py index 57a7f4af..3bbaebca 100644 --- a/bec_widgets/scan2d_plot.py +++ b/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 diff --git a/bec_widgets/scan_plot.py b/bec_widgets/scan_plot.py index b2748343..a6c4d071 100644 --- a/bec_widgets/scan_plot.py +++ b/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 diff --git a/tests/test_scan_plot.py b/tests/test_scan_plot.py index 89802a36..321304e9 100644 --- a/tests/test_scan_plot.py +++ b/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": {