Merge pull request slaclab#1042 from flowln/improve_plots_subclassing

Simplify using CurveItem subclasses in Plot subclasses
nstelter-slac · Oct 26, 2023 · e39a66e · e39a66e
2 parents 5a653c8 + 6aff422
commit e39a66e
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 36 deletions.
diff --git a/pydm/tests/widgets/test_timeplot.py b/pydm/tests/widgets/test_timeplot.py
@@ -221,6 +221,16 @@ def test_pydmtimeplot_construct(qtbot):
     assert pydm_timeplot._bottom_axis.orientation == "bottom"
 
 
+def test_pydmtimeplot_add_curve(qtbot):
+    pydm_timeplot = PyDMTimePlot()
+    qtbot.addWidget(pydm_timeplot)
+
+    curve_pv = "loc://test:timeplot:pv"
+    pydm_timeplot.addYChannel(curve_pv)
+
+    assert pydm_timeplot.findCurve(curve_pv) is not None
+
+
 @mock.patch("pydm.widgets.timeplot.TimePlotCurveItem.setData")
 @mock.patch("pyqtgraph.BarGraphItem.setOpts")
 def test_redraw_plot(mocked_set_opts, mocked_set_data, qtbot, monkeypatch):

diff --git a/pydm/widgets/archiver_time_plot.py b/pydm/widgets/archiver_time_plot.py
@@ -2,7 +2,7 @@
 import time
 import numpy as np
 from collections import OrderedDict
-from typing import List, Optional, Union
+from typing import List, Optional
 from pyqtgraph import DateAxisItem, ErrorBarItem
 from pydm.widgets.channel import PyDMChannel
 from pydm.widgets.timeplot import TimePlotCurveItem
@@ -332,26 +332,9 @@ def getArchiveBufferSize(self) -> int:
             return DEFAULT_ARCHIVE_BUFFER_SIZE
         return self._curves[0].getArchiveBufferSize()
 
-    def createCurveItem(
-        self,
-        y_channel: str,
-        plot_by_timestamps: bool,
-        name: str,
-        color: Union[QColor, str],
-        yAxisName: str,
-        useArchiveData: bool,
-        **plot_opts
-    ) -> ArchivePlotCurveItem:
+    def createCurveItem(self, *args, **kwargs) -> ArchivePlotCurveItem:
         """Create and return a curve item to be plotted"""
-        curve_item = ArchivePlotCurveItem(
-            y_channel,
-            use_archive_data=useArchiveData,
-            plot_by_timestamps=plot_by_timestamps,
-            name=name,
-            color=color,
-            yAxisName=yAxisName,
-            **plot_opts
-        )
+        curve_item = ArchivePlotCurveItem(*args, **kwargs)
         curve_item.archive_data_received_signal.connect(self.archive_data_received)
         return curve_item
 

diff --git a/pydm/widgets/eventplot.py b/pydm/widgets/eventplot.py
@@ -343,7 +343,7 @@ def addChannel(
             plot_opts["lineStyle"] = lineStyle
         if lineWidth is not None:
             plot_opts["lineWidth"] = lineWidth
-        curve = EventPlotCurveItem(
+        curve = self.createCurveItem(
             addr=channel,
             y_idx=y_idx,
             x_idx=x_idx,
@@ -359,6 +359,9 @@ def addChannel(
         self.addCurve(curve, curve_color=color, y_axis_name=yAxisName)
         curve.data_changed.connect(self.set_needs_redraw)
 
+    def createCurveItem(self, *args, **kwargs):
+        return EventPlotCurveItem(*args, *kwargs)
+
     def removeChannel(self, curve):
         """
         Remove a curve from the plot.

diff --git a/pydm/widgets/scatterplot.py b/pydm/widgets/scatterplot.py
@@ -433,7 +433,7 @@ def addChannel(
             plot_opts["lineWidth"] = lineWidth
         if redraw_mode is not None:
             plot_opts["redraw_mode"] = redraw_mode
-        curve = ScatterPlotCurveItem(
+        curve = self.createCurveItem(
             y_addr=y_channel,
             x_addr=x_channel,
             name=name,
@@ -448,6 +448,9 @@ def addChannel(
         self.addCurve(curve, curve_color=color, y_axis_name=yAxisName)
         curve.data_changed.connect(self.set_needs_redraw)
 
+    def createCurveItem(self, *args, **kwargs):
+        return ScatterPlotCurveItem(*args, **kwargs)
+
     def removeChannel(self, curve):
         """
         Remove a curve from the plot.

diff --git a/pydm/widgets/timeplot.py b/pydm/widgets/timeplot.py
@@ -519,7 +519,7 @@ def addYChannel(
 
         # Add curve
         new_curve = self.createCurveItem(
-            y_channel=y_channel,
+            channel_address=y_channel,
             plot_by_timestamps=self._plot_by_timestamps,
             plot_style=plot_style,
             name=name,
@@ -547,18 +547,8 @@ def addYChannel(
 
         return new_curve
 
-    def createCurveItem(
-        self, y_channel, plot_by_timestamps, plot_style, name, color, yAxisName, useArchiveData, **plot_opts
-    ):
-        return TimePlotCurveItem(
-            y_channel,
-            plot_by_timestamps=plot_by_timestamps,
-            plot_style=plot_style,
-            name=name,
-            color=color,
-            yAxisName=yAxisName,
-            **plot_opts
-        )
+    def createCurveItem(self, *args, **kwargs):
+        return TimePlotCurveItem(*args, **kwargs)
 
     def removeYChannel(self, curve):
         """

diff --git a/pydm/widgets/waveformplot.py b/pydm/widgets/waveformplot.py
@@ -436,7 +436,7 @@ def addChannel(
         if redraw_mode is not None:
             plot_opts["redraw_mode"] = redraw_mode
         self._needs_redraw = False
-        curve = WaveformCurveItem(
+        curve = self.createCurveItem(
             y_addr=y_channel,
             x_addr=x_channel,
             plot_style=plot_style,
@@ -457,6 +457,9 @@ def addChannel(
             curve.getViewBox().addItem(curve.bar_graph_item)
         curve.data_changed.connect(self.set_needs_redraw)
 
+    def createCurveItem(self, *args, **kwargs):
+        return WaveformCurveItem(*args, **kwargs)
+
     def removeChannel(self, curve):
         """
         Remove a curve from the plot.