engine.py

import pkgutil
import re

from nbformat import v4 as nbf
from papermill.engines import NBClientEngine, NotebookExecutionManager

_IGNORE_TAG = 'ignore'
_MANUAL_SPARK_CONFIG_FILE = 'manual_spark_config.py'


class PapermillNxcalsEngine(NBClientEngine):
    """ Papermill engine to execute notebooks requiring connection to the NxCals. The Spark session (unless already
    created) is established according to the 'manual_spark_config.py' file, additionally all cells tagged with
    'ignore' are omitted. The engine is registered as an extension to the papermill library. To use it, call :func:
    `papermill.execute_notebook` with engine_name='nxcals_engine'.
    """

    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):
        """Changes type of all cells tagged with 'ignore' to 'raw' and injects configuration of the Spark context
        to the notebook. Then calls :func: `NBClientEngine.execute_managed_notebook` to execute the parameterized
        notebook locally.
        """
        print('papermill engine')
        print(**kwargs)
        print(kwargs)
        # make all cells tagged with 'ignore' raw
        for cell in nb_man.nb.cells:
            if _IGNORE_TAG in cell.metadata.tags:
                cell.cell_type = "raw"

        # inject cell with the configuration of the Spark context
        nb_man.nb.cells.insert(0, _get_nxcals_config_cell())

        _replace_matplotlib_backend(nb_man)

        # call the papermill execution engine:
        return super().execute_managed_notebook(nb_man, kernel_name, **kwargs)
        # save_analysis_result()


def _replace_matplotlib_backend(nb_man, new_backend='inline'):
    for cell in nb_man.nb.cells:
        if cell.cell_type == 'code':
            cell.source = re.sub(r'%matplotlib \w+', '%matplotlib ' + new_backend, cell.source)


def _get_nxcals_config_cell():
    config_cell = nbf.new_code_cell(_get_nxcals_config())

    # Create papermill metadata
    config_cell.metadata.papermill = dict(
        exception=None,
        start_time=None,
        end_time=None,
        duration=None,
        status=NotebookExecutionManager.PENDING,
    )
    return config_cell


def _get_nxcals_config():
    return pkgutil.get_data('resources', _MANUAL_SPARK_CONFIG_FILE).decode("utf-8")


def save_analysis_result():
    pass