engine.py 2.28 KB
Newer Older
1
import pkgutil
2
import re
3
4
5
6

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

7
8
9
_IGNORE_TAG = 'ignore'
_MANUAL_SPARK_CONFIG_FILE = 'manual_spark_config.py'

10
11

class PapermillNxcalsEngine(NBClientEngine):
Aleksandra Mnich's avatar
Aleksandra Mnich committed
12
13
14
    """ 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:
15
16
17
18
19
    `papermill.execute_notebook` with engine_name='nxcals_engine'.
    """

    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):
Aleksandra Mnich's avatar
Aleksandra Mnich committed
20
21
22
        """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.
23
        """
Agata Malgorzata Chadaj's avatar
wip    
Agata Malgorzata Chadaj committed
24
25
26
        print('papermill engine')
        print(**kwargs)
        print(kwargs)
Aleksandra Mnich's avatar
Aleksandra Mnich committed
27
        # make all cells tagged with 'ignore' raw
28
29
30
        for cell in nb_man.nb.cells:
            if _IGNORE_TAG in cell.metadata.tags:
                cell.cell_type = "raw"
31

Aleksandra Mnich's avatar
Aleksandra Mnich committed
32
        # inject cell with the configuration of the Spark context
33
        nb_man.nb.cells.insert(0, _get_nxcals_config_cell())
34

35
36
        _replace_matplotlib_backend(nb_man)

37
38
        # call the papermill execution engine:
        return super().execute_managed_notebook(nb_man, kernel_name, **kwargs)
Agata Malgorzata Chadaj's avatar
wip    
Agata Malgorzata Chadaj committed
39
        # save_analysis_result()
40

Aleksandra Mnich's avatar
Aleksandra Mnich committed
41

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


48
49
50
51
52
53
54
55
56
57
58
59
60
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

Aleksandra Mnich's avatar
Aleksandra Mnich committed
61

62
63
def _get_nxcals_config():
    return pkgutil.get_data('resources', _MANUAL_SPARK_CONFIG_FILE).decode("utf-8")
Agata Malgorzata Chadaj's avatar
wip    
Agata Malgorzata Chadaj committed
64
65
66
67


def save_analysis_result():
    pass