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 new file mode 100644 index 0000000000000000000000000000000000000000..8f30433890c6681a8c40d6a2dff66958bb6aca18 --- /dev/null +++ b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/plot_model.py @@ -0,0 +1,33 @@ +# Mantid Repository : https://github.com/mantidproject/mantid +# +# Copyright © 2020 ISIS Rutherford Appleton Laboratory UKRI, +# NScD Oak Ridge National Laboratory, European Spallation Source +# & Institut Laue - Langevin +# SPDX - License - Identifier: GPL - 3.0 + + + +class FittingPlotModel(object): + def __init__(self): + self.plotted_workspaces = [] + + 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) + + def remove_workspace_from_plot(self, ws, ax): + if ws in self.plotted_workspaces: + self._remove_workspace_from_plot(ws, ax) + self.plotted_workspaces.remove(ws) + + @staticmethod + def _remove_workspace_from_plot(ws, ax): + 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. + self.plotted_workspaces.clear() diff --git a/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/__init__.py b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 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 new file mode 100644 index 0000000000000000000000000000000000000000..fdfc77009ad4a6f1a3d5b7cf87566df24fa07447 --- /dev/null +++ b/scripts/Engineering/gui/engineering_diffraction/tabs/fitting/plotting/test/test_plot_model.py @@ -0,0 +1,61 @@ +# Mantid Repository : https://github.com/mantidproject/mantid +# +# Copyright © 2020 ISIS Rutherford Appleton Laboratory UKRI, +# NScD Oak Ridge National Laboratory, European Spallation Source +# & Institut Laue - Langevin +# SPDX - License - Identifier: GPL - 3.0 + + +import unittest + +from mantid.py3compat import mock + +from Engineering.gui.engineering_diffraction.tabs.fitting.plotting import plot_model + +dir_path = "Engineering.gui.engineering_diffraction.tabs.fitting.plotting" + + +class FittingPlotModelTest(unittest.TestCase): + def setUp(self): + self.model = plot_model.FittingPlotModel() + + def test_adding_workspace_to_plot(self): + self.assertEqual([], 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) + 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") + ax = mock.MagicMock() + + self.model.remove_workspace_from_plot("mocked_ws", ax) + + self.assertEqual([], 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") + ax = mock.MagicMock() + + self.model.remove_workspace_from_plot("whatever", ax) + + 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"]) + 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) + + +if __name__ == '__main__': + unittest.main()