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.")