PySDM_examples.utils.notebook_vars

 1""" helper routines for use in smoke tests """
 2
 3from pathlib import Path
 4
 5import nbformat
 6
 7
 8def notebook_vars(file: Path, plot: bool):
 9    """Executes the code from all cells of the Jupyter notebook `file` and
10    returns a dictionary with the notebook variables. If the `plot` argument
11    is set to `True`, any code line within the notebook starting with `show_plot(`
12    (see [open_atmos_jupyter_utils docs](https://pypi.org/p/open_atmos_jupyter_utils))
13    is replaced with `pyplot.show() #`, otherwise it is replaced with `pyplot.gca().clear() #`
14    to match the smoke-test conventions."""
15    notebook = nbformat.read(file, nbformat.NO_CONVERT)
16    context = {}
17    for cell in notebook.cells:
18        if cell.cell_type != "markdown":
19            lines = cell.source.splitlines()
20            for i, line in enumerate(lines):
21                if line.strip().startswith("!"):
22                    lines[i] = line.replace("!", "pass #")
23                if line.strip().startswith("show_plot("):
24                    lines[i] = line.replace(
25                        "show_plot(",
26                        "from matplotlib import pyplot; "
27                        + ("pyplot.show() #" if plot else "pyplot.gca().clear() #"),
28                    )
29
30            exec("\n".join(lines), context)  # pylint: disable=exec-used
31    return context