def split_matplotlib_cells(nb):
"""
If a cell imports matplotlib, split the cell to keep the
import statement separate from the code that uses matplotlib.
This prevents a known bug in the Jupyter backend which causes
the plot object to be represented as a string instead of a canvas
when created in the cell where matplotlib is imported for the
first time (https://github.com/jupyter/notebook/issues/3523).
"""
for i in range(len(nb['cells']) - 1, -1, -1):
cell = nb['cells'][i]
if cell['cell_type'] == 'code' and 'matplotlib' in cell['source']:
code = iw.protect_ipython_magics(cell['source'])
# split cells after matplotlib imports
mapping = iw.delimit_statements(code)
tree = ast.parse(code)
visitor = iw.GetMatplotlibPyplot()
visitor.visit(tree)
if visitor.matplotlib_first:
code = iw.deprotect_ipython_magics(code)
lines = code.split('\n')
lineno_end = mapping[visitor.matplotlib_first]
split_code = '\n'.join(lines[lineno_end:]).lstrip('\n')
if split_code:
new_cell = nbformat.v4.new_code_cell(source=split_code)
nb['cells'].insert(i + 1, new_cell)
lines = lines[:lineno_end]
nb['cells'][i]['source'] = '\n'.join(lines).rstrip('\n')
def test_matplotlib_pyplot_visitor(self):
import_stmt = [
'import matplotlib',
'import matplotlib as mpl',
'import matplotlib.pyplot',
'import matplotlib.pyplot as plt',
'import matplotlib.pyplot.figure as fig',
'import ast, matplotlib',
'import ast, matplotlib as mpl',
'import ast, matplotlib.pyplot',
'import ast, matplotlib.pyplot as plt',
'import ast, matplotlib.pyplot.figure as fig',
'from matplotlib import pyplot',
'from matplotlib import pyplot as plt',
'from matplotlib.pyplot import figure',
'matplotlib.pyplot.ion()',
'matplotlib.pyplot.ioff()',
'plt.ion()',
'mpl.use("PS")',
'matplotlib.use("Agg")',
'get_ipython().run_line_magic("matplotlib", "notebook")',
]
tree = ast.parse('\n'.join(import_stmt))
v = iw.GetMatplotlibPyplot()
v.visit(tree)
# find first line where matplotlib is imported
self.assertEqual(v.matplotlib_first, 1)
# find all aliases for matplotlib
expected_mpl_aliases = ['matplotlib', 'mpl', 'matplotlib', 'mpl']
self.assertEqual(v.matplotlib_aliases, expected_mpl_aliases)
# find all aliases for matplotlib.pyplot
expected_plt_aliases = [
'matplotlib.pyplot',
'mpl.pyplot',
'matplotlib.pyplot',
'plt',
'matplotlib.pyplot',
'mpl.pyplot',
'matplotlib.pyplot',
'plt',
'pyplot',
'plt',
]
self.assertEqual(v.pyplot_aliases, expected_plt_aliases)
expected_plt_paths = {('matplotlib', 'pyplot'), ('mpl', 'pyplot'),
('plt', ), ('pyplot', )}
self.assertEqual(v.pyplot_paths, expected_plt_paths)
# find lines interactive mode, backend setup and magic functions
self.assertEqual(v.pyplot_interactive_linenos, [14, 16])
self.assertEqual(v.matplotlib_backend_linenos, [17, 18])
self.assertEqual(v.ipython_magic_linenos, [19])
cell['source'], flags=re.M)
# if matplotlib is used in this script, split cell to keep the import
# statement separate and avoid a known bug in the Jupyter backend which
# causes the plot object to be represented as a string instead of a
# canvas when created in the cell where matplotlib is imported for the
# first time (https://github.com/jupyter/notebook/issues/3523)
for i in range(len(nb['cells'])):
cell = nb['cells'][i]
if cell['cell_type'] == 'code' and 'matplotlib' in cell['source']:
code = iw.protect_ipython_magics(cell['source'])
# split cells after matplotlib imports
mapping = iw.delimit_statements(code)
tree = ast.parse(code)
visitor = iw.GetMatplotlibPyplot()
visitor.visit(tree)
if visitor.matplotlib_first:
code = iw.deprotect_ipython_magics(code)
lines = code.split('\n')
lineno_end = mapping[visitor.matplotlib_first]
split_code = '\n'.join(lines[lineno_end:]).lstrip('\n')
if split_code:
new_cell = nbformat.v4.new_code_cell(source=split_code)
nb['cells'].insert(i + 1, new_cell)
lines = lines[:lineno_end]
nb['cells'][i]['source'] = '\n'.join(lines).rstrip('\n')
break
# substitute global variables and disable OpenGL/Mayavi GUI
cell_separator = '\n##{}\n'.format(uuid.uuid4().hex)