def test_generate_script_adds_legend_command_if_legend_present(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_ax = self._gen_mock_axes(legend_=True)
mock_fig = Mock(get_axes=lambda: [mock_ax])
if hasattr(Legend, "set_draggable"):
self.assertIn('.legend().set_draggable(True)',
generate_script(mock_fig))
else:
self.assertIn('.legend().draggable()', generate_script(mock_fig))
def test_generate_script_compiles_script_correctly_with_fit(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_axis_lim_cmd, mock_autoscale_lims, mock_fit_cmds):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_axis_lim_cmd.return_value = self.ax_limit_cmds
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_fit_cmds.return_value = self.fit_cmds, self.fit_header
mock_axes1 = self._gen_mock_axes(
get_tracked_artists=lambda: [None, None],
get_lines=lambda: [None, None],
numRows=1,
numCols=2,
colNum=0)
mock_axes2 = self._gen_mock_axes(get_tracked_artists=lambda: [None],
get_lines=lambda: [None],
numRows=1,
numCols=2,
colNum=1)
mock_fig = Mock(get_axes=lambda: [mock_axes1, mock_axes2],
axes=[mock_axes1, mock_axes2])
mock_fig.canvas.manager.fit_browser.fit_result_ws_name = "OutputWorkspace"
output_script = generate_script(mock_fig, exclude_headers=True)
self.assertEqual(SAMPLE_SCRIPT_WITH_FIT, output_script)
def test_generate_tick_format_commands(self):
"""
Check that the tick format commands are correctly generated if they are set different to the default.
"""
fig, axes = plt.subplots(ncols=2,
nrows=2,
subplot_kw={'projection': 'mantid'})
for ax in fig.get_axes():
ax.plot([-10, 10], [1, 2])
# Only change the major formatter for one of the x axes and one of the y axes.
# The rest will be default, and shouldn't generate any lines in the script.
axes[0][1].xaxis.set_major_formatter(LogFormatterSciNotation())
axes[1][0].yaxis.set_major_formatter(LogFormatterSciNotation())
script = generate_script(fig)
# Check the import is there exactly once.
self.assertEqual(script.count("from matplotlib.ticker import"), 1)
# We only set the major formatter for axes[0][1].xaxis, so the command should only be present there.
self.assertNotIn("axes[0][0].xaxis.set_major_formatter", script)
self.assertIn("axes[0][1].xaxis.set_major_formatter", script)
self.assertNotIn("axes[1][0].xaxis.set_major_formatter", script)
self.assertNotIn("axes[1][1].xaxis.set_major_formatter", script)
# We only set the major formatter for axes[1][0].yaxis, so the command should only be present there.
self.assertNotIn("axes[0][0].yaxis.set_major_formatter", script)
self.assertNotIn("axes[0][1].yaxis.set_major_formatter", script)
self.assertIn("axes[1][0].yaxis.set_major_formatter", script)
self.assertNotIn("axes[1][1].yaxis.set_major_formatter", script)
def test_generate_script_compiles_script_correctly(self, mock_subplots_cmd,
mock_plot_cmd,
mock_retrieval_cmd,
mock_axis_lim_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_axis_lim_cmd.return_value = self.ax_limit_cmds
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_axes1 = self._gen_mock_axes(
get_tracked_artists=lambda: [None, None],
get_lines=lambda: [None, None],
numRows=1,
numCols=2,
colNum=0)
mock_axes2 = self._gen_mock_axes(get_tracked_artists=lambda: [None],
get_lines=lambda: [None],
numRows=1,
numCols=2,
colNum=1)
mock_fig = Mock(get_axes=lambda: [mock_axes1, mock_axes2])
output_script = generate_script(mock_fig, exclude_headers=True)
self.assertEqual(SAMPLE_SCRIPT, output_script)
def test_generate_script_adds_minor_and_major_tick_kw(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_ax = self._gen_mock_axes()
mock_ax.xaxis.minor.locator = Mock()
mock_ax.xaxis.major.locator = Mock()
mock_ax.xaxis.majorTicks = [None]
mock_ax.xaxis.minorTicks = [None]
# If this fails then the internals of matplotlib have changed, and axes.py for ticks needs to be re-thought.
# This fails because it is an object being made to the spec defined in matplotlib so should definetely contain
# This object, mock or not.
self.assertTrue(hasattr(mock_ax.xaxis, "_major_tick_kw"))
self.assertTrue(hasattr(mock_ax.xaxis, "_minor_tick_kw"))
mock_minor_kw = "{gridOn: False, show: True, width: 6, length: 10}"
mock_major_kw = "{gridOn: True, show: True, width: 20, length: 15}"
mock_ax.xaxis._major_tick_kw = mock_major_kw
mock_ax.xaxis._minor_tick_kw = mock_minor_kw
mock_fig = Mock(get_axes=lambda: [mock_ax])
mock_fig.canvas.manager.fit_browser.fit_result_ws_name = ""
commands = generate_script(mock_fig)
self.assertIn(mock_major_kw, commands)
self.assertIn(mock_minor_kw, commands)
def test_generate_script_does_not_add_legend_command_if_figure_has_no_legend(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_fig = Mock(get_axes=lambda: [self._gen_mock_axes()])
self.assertNotIn('.legend()', generate_script(mock_fig))
def generate_plot_script_file(self):
script = generate_script(self.canvas.figure)
filepath = open_a_file_dialog(parent=self.canvas,
default_suffix=".py",
file_filter="Python Files (*.py)",
accept_mode=QFileDialog.AcceptSave,
file_mode=QFileDialog.AnyFile)
if filepath:
try:
with open(filepath, 'w') as f:
f.write(script)
except IOError as io_error:
logger.error("Could not write file: {}\n{}" "".format(filepath, io_error))
def test_generate_script_adds_minor_ticks_command_if_axes_has_minor_ticks(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_ax = self._gen_mock_axes()
mock_ax.xaxis.minor.locator = Mock()
mock_fig = Mock(get_axes=lambda: [mock_ax])
mock_fig.canvas.manager.fit_browser.fit_result_ws_name = ""
self.assertIn('axes.minorticks_on()', generate_script(mock_fig))
def test_generate_tick_format_commands_log_scale(self):
"""
Check that tick format commands are correctly generated for a log axis scale.
"""
fig = plt.figure()
axes = fig.add_subplot(1, 1, 1, projection='mantid')
axes.plot([-10, 10], [1, 2])
axes.set_yscale('log')
axes.yaxis.set_major_formatter(ScalarFormatter())
script = generate_script(fig)
# Scalar format is not the default for log axis, so should appear in script.
self.assertIn("axes.yaxis.set_major_formatter(ScalarFormatter(",
script)
# LogFormatterSciNotation is default for log axis, so shouldn't appear in script for minor or major ticks.
self.assertNotIn("_formatter(LogFormatterSciNotation", script)
def test_generate_script_adds_grid_command_if_axes_has_grid(
self, mock_subplots_cmd, mock_plot_cmd, mock_retrieval_cmd,
mock_autoscale_lims):
mock_retrieval_cmd.return_value = self.retrieval_cmds
mock_subplots_cmd.return_value = self.subplots_cmd
mock_plot_cmd.return_value = self.plot_cmd
mock_autoscale_lims.return_value = (-0.02, 1.02)
mock_ax = self._gen_mock_axes()
mock_ax.show_minor_gridlines = True
mock_ax.xaxis._gridOnMajor = True
mock_ax.yaxis._gridOnMajor = True
mock_fig = Mock(get_axes=lambda: [mock_ax])
mock_fig.canvas.manager.fit_browser.fit_result_ws_name = ""
self.assertIn('axes.grid(True, axis=\'both\', which=\'both\')',
generate_script(mock_fig))
def test_generate_tick_commands_for_tiled_plot(self):
"""
Check that the tick commands are generated for every plot in the figure.
"""
fig, axes = plt.subplots(ncols=2,
nrows=2,
subplot_kw={'projection': 'mantid'})
for ax in fig.get_axes():
ax.plot([-10, 10], [1, 2])
script = generate_script(fig)
# Should be axes[i][j].tick_params for multiple subplots.
self.assertNotIn("axes.tick_params", script)
self.assertIn("axes[0][0].tick_params", script)
self.assertIn("axes[0][1].tick_params", script)
self.assertIn("axes[1][0].tick_params", script)
self.assertIn("axes[1][1].tick_params", script)
plt.close()
del fig
def test_facecolor_commands_for_tiled_plot(self):
"""
Set a canvas colour (facecolor) for only one of the subplots in a tiled plot and check
that commands are generated for only that one.
"""
fig, axes = plt.subplots(ncols=2,
nrows=2,
subplot_kw={'projection': 'mantid'})
for ax in fig.get_axes():
ax.plot([-10, 10], [1, 2])
# Only set the facecolor on one subplot.
axes[1][0].set_facecolor('#c0ffee')
script = generate_script(fig)
# Should be axes[i][j].set_facecolor for multiple subplots.
self.assertNotIn("axes.set_facecolor", script)
# These subplots have default canvas colour so should've have the facecolor command.
self.assertNotIn("axes[0][0].set_facecolor", script)
self.assertNotIn("axes[0][1].set_facecolor", script)
self.assertNotIn("axes[1][1].set_facecolor", script)
# This subplot should have our new colour.
self.assertIn("axes[1][0].set_facecolor('#c0ffee')", script)
plt.close()
del fig
def test_generate_script_returns_None_if_no_MantidAxes_in_figure(self):
mock_fig = Mock(get_axes=lambda: [Mock(spec=Axes)])
self.assertEqual(None, generate_script(mock_fig))
def generate_plot_script_clipboard(self):
script = generate_script(self.canvas.figure, exclude_headers=True)
QApplication.clipboard().setText(script)
logger.notice("Plotting script copied to clipboard.")
