A mocked out version of MatPlotLib

mockplt.py

class MockPlt:
    '''
    Mock MatPlotLib library that can be used to capture plot data.
    '''
    def __init__(self):
        self.plots = []
        self._reset_plot()
    def show(self, **kwargs):
        self.plots.append(self.active_plot)
        self._reset_plot()
    def unshown_plots(self):
        return self.active_plot['data']
    def __repr__(self):
        return repr(self.plots)
    def __str__(self):
        return str(self.plots)
    def _reset_plot(self):
        self.active_plot = {'data': [], 
                            'xlabel': None, 'ylabel': None, 
                            'title': None, 'legend': False}
    def hist(self, data, **kwargs):
        histogram = {'type': 'Histogram', 'values': data}
        if 'label' in kwargs and kwargs['label']:
            histogram['label'] = kwargs['label']
        self.active_plot['data'].append(histogram)
    def plot(self, xs, ys=None, style=None, **kwargs):
        line_plot = {'type': 'Line'}
        if ys == None or isinstance(ys, str):
            line_plot['x'] = range(len(xs))
            line_plot['y'] = xs
        else:
            line_plot['x'] = xs
            line_plot['y'] = ys
        if 'label' in kwargs and kwargs['label']:
            line_plot['label'] = kwargs['label']
        self.active_plot['data'].append(line_plot)
    def scatter(self, xs, ys, style=None, **kwargs):
        scatter_plot = {'type': 'Scatter', 'x': xs, 'y': ys}
        if 'label' in kwargs and kwargs['label']:
            scatter_plot['label'] = kwargs['label']
        self.active_plot['data'].append(scatter_plot)
    def xlabel(self, label, **kwargs):
        self.active_plot['xlabel'] = label
    def title(self, label, **kwargs):
        self.active_plot['title'] = label
    def ylabel(self, label, **kwargs):
        self.active_plot['ylabel'] = label
    def legend(self, *args, **kwargs):
        if len(args) == 2 and isinstance(args[1], (list, tuple)):
            for label, plot in zip(args[1], self.active_plot['data']):
                plot['label'] = label
        elif len(args) == 1 and isinstance(args[0], (list, tuple)):
            for label, plot in zip(args[0], self.active_plot['data']):
                plot['label'] = label
        self.active_plot['legend'] = True
    def _generate_patches(self):
        def dummy(*pos, **kwargs):
            pass
        return dict(hist=self.hist, plot=self.plot, 
                    scatter=self.scatter, show=self.show,
                    xlabel=self.xlabel, ylabel=self.ylabel, 
                    title=self.title, legend=self.legend,
                    xticks=dummy, yticks=dummy,
                    autoscale=dummy, axhline=dummy,
                    axhspan=dummy, axvline=dummy,
                    axvspan=dummy, clf=dummy,
                    cla=dummy, close=dummy,
                    figlegend=dummy, figimage=dummy,
                    suptitle=dummy, text=dummy,
                    figure=dummy,
                    tick_params=dummy, ticklabel_format=dummy,
                    tight_layout=dummy, xkcd=dummy,
                    xlim=dummy, ylim=dummy,
                    xscale=dummy, yscale=dummy)
      
class MainTest(unittest.TestCase):
    def test_a_function(self):
        mock_plt = MockPlt()
        plt_patches = mock_plt._generate_patches()
        with patch.multiple('matplotlib.pyplot', **plt_patches):
            # run student code
        # Test if right number of plots
        plots = mock_plt.plots
        self.assertEqual(1, len(plots), msg="Too many graphs")
        # Test if histogram has data
        histogram = plots[0]
        self.assertTrue(histogram['data'], msg="Your histogram is empty.")
        # Test if histogram
        data = histogram['data']
        self.assertEqual("Histogram", data[0]['type'], msg="Your graph is not a histogram")