File: parameterize.py

package info (click to toggle)
python-papermill 2.6.0-3.1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 2,216 kB
  • sloc: python: 4,977; makefile: 17; sh: 5
file content (121 lines) | stat: -rw-r--r-- 3,901 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
from datetime import datetime
from uuid import uuid4

import nbformat

from .engines import papermill_engines
from .exceptions import PapermillMissingParameterException
from .iorw import read_yaml_file
from .log import logger
from .translators import translate_parameters
from .utils import find_first_tagged_cell_index


def add_builtin_parameters(parameters):
    """Add built-in parameters to a dictionary of parameters

    Parameters
    ----------
    parameters : dict
       Dictionary of parameters provided by the user
    """
    with_builtin_parameters = {
        "pm": {
            "run_uuid": str(uuid4()),
            "current_datetime_local": datetime.now(),
            "current_datetime_utc": datetime.utcnow(),
        }
    }

    if parameters is not None:
        with_builtin_parameters.update(parameters)

    return with_builtin_parameters


def parameterize_path(path, parameters):
    """Format a path with a provided dictionary of parameters

    Parameters
    ----------
    path : string or nbformat.NotebookNode or None
       Path with optional parameters, as a python format string. If path is a NotebookNode
       or None, the path is returned without modification
    parameters : dict or None
       Arbitrary keyword arguments to fill in the path
    """
    if path is None or isinstance(path, nbformat.NotebookNode):
        return path

    if parameters is None:
        parameters = {}

    try:
        return path.format(**parameters)
    except KeyError as key_error:
        raise PapermillMissingParameterException(f"Missing parameter {key_error}")


def parameterize_notebook(
    nb,
    parameters,
    report_mode=False,
    comment='Parameters',
    kernel_name=None,
    language=None,
    engine_name=None,
):
    """Assigned parameters into the appropriate place in the input notebook

    Parameters
    ----------
    nb : NotebookNode
       Executable notebook object
    parameters : dict
       Arbitrary keyword arguments to pass as notebook parameters
    report_mode : bool, optional
       Flag to set report mode
    comment : str, optional
        Comment added to the injected cell
    """
    # Load from a file if 'parameters' is a string.
    if isinstance(parameters, str):
        parameters = read_yaml_file(parameters)

    # Fetch out the name and language from the notebook document by dropping-down into the engine's implementation
    kernel_name = papermill_engines.nb_kernel_name(engine_name, nb, kernel_name)
    language = papermill_engines.nb_language(engine_name, nb, language)

    # Generate parameter content based on the kernel_name
    param_content = translate_parameters(kernel_name, language, parameters, comment)

    # Upgrade the Notebook to the latest v4 before writing into it
    nb = nbformat.v4.upgrade(nb)

    newcell = nbformat.v4.new_code_cell(source=param_content)
    newcell.metadata['tags'] = ['injected-parameters']

    if report_mode:
        newcell.metadata['jupyter'] = newcell.get('jupyter', {})
        newcell.metadata['jupyter']['source_hidden'] = True

    param_cell_index = find_first_tagged_cell_index(nb, 'parameters')
    injected_cell_index = find_first_tagged_cell_index(nb, 'injected-parameters')
    if injected_cell_index >= 0:
        # Replace the injected cell with a new version
        before = nb.cells[:injected_cell_index]
        after = nb.cells[injected_cell_index + 1 :]
    elif param_cell_index >= 0:
        # Add an injected cell after the parameter cell
        before = nb.cells[: param_cell_index + 1]
        after = nb.cells[param_cell_index + 1 :]
    else:
        # Inject to the top of the notebook
        logger.warning("Input notebook does not contain a cell with tag 'parameters'")
        before = []
        after = nb.cells

    nb.cells = before + [newcell] + after
    nb.metadata.papermill['parameters'] = parameters

    return nb