"""Tests for pylab tools module. """ # Copyright (c) IPython Development Team. # Distributed under the terms of the Modified BSD License. from binascii import a2b_base64 from io import BytesIO import matplotlib matplotlib.use('Agg') from matplotlib.figure import Figure from nose import SkipTest import nose.tools as nt from matplotlib import pyplot as plt import matplotlib_inline import numpy as np from IPython.core.getipython import get_ipython from IPython.core.interactiveshell import InteractiveShell from IPython.core.display import _PNG, _JPEG from .. import pylabtools as pt from IPython.testing import decorators as dec def test_figure_to_svg(): # simple empty-figure test fig = plt.figure() nt.assert_equal(pt.print_figure(fig, 'svg'), None) plt.close('all') # simple check for at least svg-looking output fig = plt.figure() ax = fig.add_subplot(1,1,1) ax.plot([1,2,3]) plt.draw() svg = pt.print_figure(fig, 'svg')[:100].lower() nt.assert_in(u'doctype svg', svg) def _check_pil_jpeg_bytes(): """Skip if PIL can't write JPEGs to BytesIO objects""" # PIL's JPEG plugin can't write to BytesIO objects # Pillow fixes this from PIL import Image buf = BytesIO() img = Image.new("RGB", (4,4)) try: img.save(buf, 'jpeg') except Exception as e: ename = e.__class__.__name__ raise SkipTest("PIL can't write JPEG to BytesIO: %s: %s" % (ename, e)) from e @dec.skip_without("PIL.Image") def test_figure_to_jpeg(): _check_pil_jpeg_bytes() # simple check for at least jpeg-looking output fig = plt.figure() ax = fig.add_subplot(1,1,1) ax.plot([1,2,3]) plt.draw() jpeg = pt.print_figure(fig, 'jpeg', pil_kwargs={'optimize': 50})[:100].lower() assert jpeg.startswith(_JPEG) def test_retina_figure(): # simple empty-figure test fig = plt.figure() nt.assert_equal(pt.retina_figure(fig), None) plt.close('all') fig = plt.figure() ax = fig.add_subplot(1,1,1) ax.plot([1,2,3]) plt.draw() png, md = pt.retina_figure(fig) assert png.startswith(_PNG) nt.assert_in('width', md) nt.assert_in('height', md) _fmt_mime_map = { 'png': 'image/png', 'jpeg': 'image/jpeg', 'pdf': 'application/pdf', 'retina': 'image/png', 'svg': 'image/svg+xml', } def test_select_figure_formats_str(): ip = get_ipython() for fmt, active_mime in _fmt_mime_map.items(): pt.select_figure_formats(ip, fmt) for mime, f in ip.display_formatter.formatters.items(): if mime == active_mime: nt.assert_in(Figure, f) else: nt.assert_not_in(Figure, f) def test_select_figure_formats_kwargs(): ip = get_ipython() kwargs = dict(quality=10, bbox_inches='tight') pt.select_figure_formats(ip, 'png', **kwargs) formatter = ip.display_formatter.formatters['image/png'] f = formatter.lookup_by_type(Figure) cell = f.keywords expected = kwargs expected["base64"] = True expected["fmt"] = "png" assert cell == expected # check that the formatter doesn't raise fig = plt.figure() ax = fig.add_subplot(1,1,1) ax.plot([1,2,3]) plt.draw() formatter.enabled = True png = formatter(fig) assert isinstance(png, str) png_bytes = a2b_base64(png) assert png_bytes.startswith(_PNG) def test_select_figure_formats_set(): ip = get_ipython() for fmts in [ {'png', 'svg'}, ['png'], ('jpeg', 'pdf', 'retina'), {'svg'}, ]: active_mimes = {_fmt_mime_map[fmt] for fmt in fmts} pt.select_figure_formats(ip, fmts) for mime, f in ip.display_formatter.formatters.items(): if mime in active_mimes: nt.assert_in(Figure, f) else: nt.assert_not_in(Figure, f) def test_select_figure_formats_bad(): ip = get_ipython() with nt.assert_raises(ValueError): pt.select_figure_formats(ip, 'foo') with nt.assert_raises(ValueError): pt.select_figure_formats(ip, {'png', 'foo'}) with nt.assert_raises(ValueError): pt.select_figure_formats(ip, ['retina', 'pdf', 'bar', 'bad']) def test_import_pylab(): ns = {} pt.import_pylab(ns, import_all=False) nt.assert_true('plt' in ns) nt.assert_equal(ns['np'], np) from traitlets.config import Config class TestPylabSwitch(object): class Shell(InteractiveShell): def init_history(self): """Sets up the command history, and starts regular autosaves.""" self.config.HistoryManager.hist_file = ":memory:" super().init_history() def enable_gui(self, gui): pass def setup(self): import matplotlib def act_mpl(backend): matplotlib.rcParams['backend'] = backend # Save rcParams since they get modified self._saved_rcParams = matplotlib.rcParams self._saved_rcParamsOrig = matplotlib.rcParamsOrig matplotlib.rcParams = dict(backend='Qt4Agg') matplotlib.rcParamsOrig = dict(backend='Qt4Agg') # Mock out functions self._save_am = pt.activate_matplotlib pt.activate_matplotlib = act_mpl self._save_ip = pt.import_pylab pt.import_pylab = lambda *a,**kw:None self._save_cis = matplotlib_inline.backend_inline.configure_inline_support matplotlib_inline.backend_inline.configure_inline_support = ( lambda *a, **kw: None ) def teardown(self): pt.activate_matplotlib = self._save_am pt.import_pylab = self._save_ip matplotlib_inline.backend_inline.configure_inline_support = self._save_cis import matplotlib matplotlib.rcParams = self._saved_rcParams matplotlib.rcParamsOrig = self._saved_rcParamsOrig def test_qt(self): s = self.Shell() gui, backend = s.enable_matplotlib(None) nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') gui, backend = s.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') nt.assert_equal(s.pylab_gui_select, 'qt') gui, backend = s.enable_matplotlib('qt') nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') gui, backend = s.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') nt.assert_equal(s.pylab_gui_select, 'qt') gui, backend = s.enable_matplotlib() nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') def test_inline(self): s = self.Shell() gui, backend = s.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') nt.assert_equal(s.pylab_gui_select, None) gui, backend = s.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') nt.assert_equal(s.pylab_gui_select, None) gui, backend = s.enable_matplotlib('qt') nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') def test_inline_twice(self): "Using '%matplotlib inline' twice should not reset formatters" ip = self.Shell() gui, backend = ip.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') fmts = {'png'} active_mimes = {_fmt_mime_map[fmt] for fmt in fmts} pt.select_figure_formats(ip, fmts) gui, backend = ip.enable_matplotlib('inline') nt.assert_equal(gui, 'inline') for mime, f in ip.display_formatter.formatters.items(): if mime in active_mimes: nt.assert_in(Figure, f) else: nt.assert_not_in(Figure, f) def test_qt_gtk(self): s = self.Shell() gui, backend = s.enable_matplotlib('qt') nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') gui, backend = s.enable_matplotlib('gtk') nt.assert_equal(gui, 'qt') nt.assert_equal(s.pylab_gui_select, 'qt') def test_no_gui_backends(): for k in ['agg', 'svg', 'pdf', 'ps']: assert k not in pt.backend2gui def test_figure_no_canvas(): fig = Figure() fig.canvas = None pt.print_figure(fig)