RE #28251 Refactor removing artists for waterfall plots

Framework/PythonInterface/mantid/plots/mantidaxes.py

@@ -13,6 +13,8 @@ except ImportError:
 import copy
 import numpy as np
 
+from copy import copy
+
 from matplotlib.axes import Axes
 from matplotlib.collections import Collection, PolyCollection
 from matplotlib.colors import Colormap
@@ -285,6 +287,51 @@ class MantidAxes(Axes):
         return tracked_artists
 
     def remove_workspace_artists(self, workspace):
+        if self.is_waterfall():
+            return self._remove_workspace_artists_waterfall(workspace=workspace)
+        else:
+            return self._remove_workspace_artists(workspace)
+
+    def remove_artists_if(self, unary_predicate):
+        if self.is_waterfall():
+            return self._remove_workspace_artists_waterfall(predicate=unary_predicate)
+        else:
+            return self._remove_artists_if(unary_predicate)
+
+    def _remove_workspace_artists_waterfall(self, workspace=None, predicate=None):
+        """
+        Perform the steps necessary to maintain waterfall plot settings before removing
+        the artists. Output is based on the inner function.
+        If workspace is set, uses _remove_workspace_artists()
+        otherwise if predicate is set, uses _remove_artists_if()
+        otherwise raises a RuntimeError.
+        :param workspace: A Workspace object
+        :param predicate: A unary predicate used to select artists.
+        :return: The output of the inner function.
+        """
+        waterfall_x_offset = copy(self.waterfall_x_offset)
+        waterfall_y_offset = copy(self.waterfall_y_offset)
+        has_fills = self.waterfall_has_fill()
+
+        self.update_waterfall(0, 0)
+
+        if workspace is not None:
+            output = self._remove_workspace_artists(workspace)
+        elif predicate is not None:
+            output = self._remove_artists_if(predicate)
+        else:
+            raise RuntimeError("A workspace or predicate is required.")
+
+        self.update_waterfall(waterfall_x_offset, waterfall_y_offset)
+
+        if len(self.lines) == 1:  # Can't have waterfall plots with only one line.
+            self.set_waterfall(False)
+        elif has_fills:
+            datafunctions.waterfall_update_fill(self)
+
+        return output
+
+    def _remove_workspace_artists(self, workspace):
         """
         Remove the artists reference by this workspace (if any) and return True
         if the axes is then empty
@@ -302,7 +349,7 @@ class MantidAxes(Axes):
 
         return True
 
-    def remove_artists_if(self, unary_predicate):
+    def _remove_artists_if(self, unary_predicate):
         """
         Remove any artists which satisfy the predicate and return True
         if the axes is then empty

qt/applications/workbench/workbench/plotting/figuremanager.py

@@ -11,7 +11,6 @@
 import copy
 import sys
 from functools import wraps
-
 import matplotlib
 from matplotlib._pylab_helpers import Gcf
 from matplotlib.axes import Axes
@@ -96,12 +95,6 @@ class FigureManagerADSObserver(AnalysisDataServiceObserver):
         for ax in all_axes:
             if isinstance(ax, MantidAxes):
                 to_redraw = ax.remove_workspace_artists(workspace)
-                if ax.is_waterfall():
-                    if len(ax.lines) == 1:  # Can't have waterfall plots with only one line.
-                        ax.set_waterfall(False)
-                    else:
-                        if ax.waterfall_has_fill():
-                            datafunctions.waterfall_update_fill(ax)
             else:
                 to_redraw = False
             # We check for images and lines below as a pseudo check for an axes being

qt/python/mantidqt/widgets/plotconfigdialog/curvestabwidget/presenter.py

@@ -6,7 +6,6 @@
 # SPDX - License - Identifier: GPL - 3.0 +
 #  This file is part of the mantid workbench.
 
-from matplotlib.collections import PolyCollection
 from matplotlib.lines import Line2D
 
 from mantid.plots.legend import LegendProperties
@@ -187,36 +186,11 @@ class CurvesTabWidgetPresenter:
         if ax.legend_:
             self.legend_props = LegendProperties.from_legend(ax.legend_)
 
-        waterfall = False
-        if isinstance(ax, MantidAxes):
-            waterfall = ax.is_waterfall()
-
-        if waterfall:
-            # Waterfall plots are reset so they can be reconverted after the curve is removed.
-            x, y = ax.waterfall_x_offset, ax.waterfall_y_offset
-            ax.update_waterfall(0, 0)
-
-            # If the curves have a fill, the one which corresponds to the curve being removed also needs to be removed.
-            current_curve_index = self.view.select_curve_combo_box.currentIndex()
-            i = 0
-            for collection in ax.collections:
-                if isinstance(collection, PolyCollection):
-                    if current_curve_index == i:
-                        ax.collections.remove(collection)
-                        break
-                    i = i + 1
-
         # Remove curve from ax and remove from curve names dictionary
         remove_curve_from_ax(self.get_selected_curve())
         self.curve_names_dict.pop(self.view.get_selected_curve_name())
         self.set_apply_to_all_buttons_enabled()
 
-        # If there is now only one curve on a waterfall plot, the plot becomes non-waterfall.
-        if waterfall:
-            ax.update_waterfall(x, y)
-            if len(ax.get_lines()) <= 1:
-                ax.set_waterfall(False)
-
         ax = self.get_selected_ax()
         # Update the legend and redraw
         FigureErrorsManager.update_limits_and_legend(ax, self.legend_props)