Use common code for names to workspace conversion.

Fixes 2D plotting for WorkspaceGroups.
Refs #21567

qt/applications/workbench/workbench/plugins/workspacewidget.py

@@ -20,7 +20,8 @@ from __future__ import (absolute_import, unicode_literals)
 import functools
 
 # third-party library imports
-from mantid.api import AnalysisDataService
+from mantid.api import AnalysisDataService, MatrixWorkspace, WorkspaceGroup
+from mantid.kernel import Logger
 from mantidqt.widgets.workspacewidget.mantidtreemodel import MantidTreeModel
 from mantidqt.widgets.workspacewidget.plotselectiondialog import get_plot_selection
 from mantidqt.widgets.workspacewidget.workspacetreewidget import WorkspaceTreeWidget
@@ -31,13 +32,37 @@ from workbench.plugins.base import PluginWidget
 from workbench.plotting.functions import pcolormesh, plot
 
 
+LOGGER = Logger(b"workspacewidget")
+
+
 def _workspaces_from_names(names):
-    """Convert a list of workspace names to a list of workspace handles
+    """
+    Convert a list of workspace names to a list of MatrixWorkspace handles. Any WorkspaceGroup
+    encountered is flattened and its members inserted into the list at this point
 
-    :param names: List of names of workspaces
+    Flattens any workspace groups with the list, preserving the order of the remaining elements
+    :param names: A list of workspace names
+    :return: A list where each element is a single MatrixWorkspace
     """
     ads = AnalysisDataService.Instance()
-    return [ads[name.encode('utf-8')] for name in names]
+    flat = []
+    for name in names:
+        try:
+            ws = ads[name.encode('utf-8')]
+        except KeyError:
+            LOGGER.warning("Skipping {} as it does not exist".format(name))
+            continue
+
+        if isinstance(ws, MatrixWorkspace):
+            flat.append(ws)
+        elif isinstance(ws, WorkspaceGroup):
+            group_len = len(ws)
+            for i in range(group_len):
+                flat.append(ws[i])
+        else:
+            LOGGER.warning("{} it is not a MatrixWorkspace or WorkspaceGroup.".format(name))
+
+    return flat
 
 
 class WorkspaceWidget(PluginWidget):

qt/python/mantidqt/widgets/workspacewidget/plotselectiondialog.py

@@ -19,8 +19,6 @@ from __future__ import (absolute_import, unicode_literals)
 # std imports
 
 # 3rd party imports
-from mantid.api import MatrixWorkspace, WorkspaceGroup
-from mantid.kernel import Logger
 import qtawesome as qta
 from qtpy.QtWidgets import QDialogButtonBox
 
@@ -32,9 +30,6 @@ RANGE_SPECIFIER = '-'
 PLACEHOLDER_FORMAT = 'valid range: {}' + RANGE_SPECIFIER + '{}'
 RED_ASTERISK = None
 
-# Globals
-LOGGER = Logger(b"plotselectiondialog")
-
 
 def red_asterisk():
     global RED_ASTERISK
@@ -172,21 +167,20 @@ def get_plot_selection(workspaces, parent_widget):
     dialog will only be shown in the case where all workspaces
     have more than 1 spectrum
 
-    :param workspaces: A list of workspaces that will be plotted
+    :param workspaces: A list of MatrixWorkspaces that will be plotted
     :param parent_widget: A parent_widget to use for the input selection dialog
     :returns: Either a PlotSelection object containing the details of workspaces to plot or None indicating
     the request was cancelled
     """
-    flat_list = flatten(workspaces)
-    single_spectra_ws = [wksp.getNumberHistograms() for wksp in flat_list if wksp.getNumberHistograms() == 1]
+    single_spectra_ws = [wksp.getNumberHistograms() for wksp in workspaces if wksp.getNumberHistograms() == 1]
     if len(single_spectra_ws) > 0:
         # At least 1 workspace contains only a single spectrum so this is all
         # that is possible to plot for all of them
-        selection = PlotSelection(flat_list)
+        selection = PlotSelection(workspaces)
         selection.wksp_indices = [0]
         return selection
     else:
-        selection_dlg = PlotSelectionDialog(flat_list, parent=parent_widget)
+        selection_dlg = PlotSelectionDialog(workspaces, parent=parent_widget)
         res = selection_dlg.exec_()
         if res == PlotSelectionDialog.Rejected:
             # cancelled
@@ -198,26 +192,6 @@ def get_plot_selection(workspaces, parent_widget):
             return user_selection
 
 
-def flatten(workspaces):
-    """
-    Flattens any workspace groups with the list, preserving the order of the remaining elements
-    :param workspaces: A list containing workspace handles
-    :return: A list where each element is a single MatrixWorkspace
-    """
-    flat = []
-    for ws in workspaces:
-        if isinstance(ws, MatrixWorkspace):
-            flat.append(ws)
-        elif isinstance(ws, WorkspaceGroup):
-            group_len = len(ws)
-            for i in range(group_len):
-                flat.append(ws[i])
-        else:
-            LOGGER.warning("Attempting to plot item that is not a MatrixWorkspace or WorkspaceGroup")
-
-    return flat
-
-
 def parse_selection_str(txt, min_val, max_val):
     """Parse an input string containing plot index selection.