diff --git a/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py index 8f30433890c6681a8c40d6a2dff66958bb6aca18..4f8ae77586e94c674b77c99c810cedbfe011f787 100644 --- a/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py +++ b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py @@ -8,14 +8,14 @@ class FittingPlotModel(object): def __init__(self): - self.plotted_workspaces = [] + self.plotted_workspaces = set() def get_plotted_workspaces(self): return self.plotted_workspaces def add_workspace_to_plot(self, ws, ax, plot_kwargs): ax.plot(ws, **plot_kwargs) - self.plotted_workspaces.append(ws) + self.plotted_workspaces.add(ws) def remove_workspace_from_plot(self, ws, ax): if ws in self.plotted_workspaces: @@ -27,7 +27,5 @@ class FittingPlotModel(object): ax.remove_workspace_artists(ws) def remove_all_workspaces_from_plot(self, ax): - for ws in self.plotted_workspaces: - self._remove_workspace_from_plot(ws, ax) - # Should already be empty, but just to be sure. + ax.cla() self.plotted_workspaces.clear() diff --git a/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py index fdfc77009ad4a6f1a3d5b7cf87566df24fa07447..32c7aa275bce61ea993bf7b29365c8a42fa89711 100644 --- a/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py +++ b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py @@ -19,42 +19,40 @@ class FittingPlotModelTest(unittest.TestCase): self.model = plot_model.FittingPlotModel() def test_adding_workspace_to_plot(self): - self.assertEqual([], self.model.plotted_workspaces) + self.assertEqual(set(), self.model.plotted_workspaces) ax = mock.MagicMock() self.model.add_workspace_to_plot("mocked_ws", ax, {"linestyle": "x"}) - self.assertEqual(["mocked_ws"], self.model.plotted_workspaces) + self.assertEqual({"mocked_ws"}, self.model.plotted_workspaces) ax.plot.assert_called_once_with("mocked_ws", linestyle="x") def test_removing_single_tracked_workspace_from_plot(self): - self.model.plotted_workspaces.append("mocked_ws") + self.model.plotted_workspaces.add("mocked_ws") ax = mock.MagicMock() self.model.remove_workspace_from_plot("mocked_ws", ax) - self.assertEqual([], self.model.plotted_workspaces) + self.assertEqual(set(), self.model.plotted_workspaces) ax.remove_workspace_artists.assert_called_once_with("mocked_ws") def test_removing_not_tracked_workspace_from_plot(self): - self.model.plotted_workspaces.append("mocked_ws") + self.model.plotted_workspaces.add("mocked_ws") ax = mock.MagicMock() self.model.remove_workspace_from_plot("whatever", ax) - self.assertEqual(["mocked_ws"], self.model.plotted_workspaces) + self.assertEqual({"mocked_ws"}, self.model.plotted_workspaces) ax.remove_workspace_artists.assert_not_called() def test_removing_all_workspaces_from_plot(self): - self.model.plotted_workspaces.extend(["mocked_ws", "mock_ws_2"]) + self.model.plotted_workspaces.update({"mocked_ws", "mock_ws_2"}) ax = mock.MagicMock() self.model.remove_all_workspaces_from_plot(ax) - self.assertEqual([], self.model.plotted_workspaces) - ax.remove_workspace_artists.assert_any_call("mocked_ws") - ax.remove_workspace_artists.assert_any_call("mock_ws_2") - self.assertEqual(2, ax.remove_workspace_artists.call_count) + self.assertEqual(set(), self.model.plotted_workspaces) + self.assertEqual(1, ax.cla.call_count) if __name__ == '__main__':