Store function on Nuke NoOp (or on nuke.root()) node and have it made available upon loading of script

storefunc.py

"""Store function on node

Usage example:

    .. code-block::

        # Control node
        ctrl_node_name = 'MY_SCRIPT_NODE'  # if None, store on nuke.root()

        # Define your functions and add them to the dictionary
        functions = {}
        function_name = 'my_function'
        function_text = 'def printer(string): print(string)'
        functions[function_name] = function_text

        # Add functions to control node, onScriptLoad
        _add_functions(functions, ctrl_node_name)

    Then in order to make the functions available without re-opening
    the script, run the following code:

    .. code_block::

        ctrl_node = _control_node(ctrl_node_name)
        for function_name, function_text in functions.items():
            exec(ctrl_node[function_name].getText())

    Execute the custom function from this example in Nuke:

    .. code-block::

        my_function('foo')

"""

import nuke


def __node_exists(node_name):
    """Checks if node name exists"""

    for nuke_node in nuke.allNodes():
        if node_name == nuke_node.name():
            return True
    return False


def _control_node(ctrl_node_name):
    """Returns the control node (and creates it if it does not exist)"""

    if ctrl_node_name is None:
        return nuke.root()
    else:
        if not __node_exists(ctrl_node_name):
            ctrl_node = nuke.createNode('NoOp')
            ctrl_node.setName(ctrl_node_name)
        else:
            ctrl_node = nuke.toNode(ctrl_node_name)
        return ctrl_node


def _add_function_to_ctrl_node(function_name, function_text, ctrl_node_name):
    """Adds given function to knob on given node"""
    ctrl_node = _control_node(ctrl_node_name)
    print('Add to ' + str(ctrl_node))

    if function_name not in ctrl_node.knobs():
        function_knob = nuke.Multiline_Eval_String_Knob(function_name)
        ctrl_node.addKnob(function_knob)
        ctrl_node.knob(function_name).setText(function_text)


def __register_onScriptLoad_prerequisites(ctrl_node_name):
    """Register prerequisites into onScriptLoad"""

    knob_text = nuke.root().knob('onScriptLoad').getValue()
    if ctrl_node_name is None:
        prereqs = 'import nuke; ctrl = nuke.root(); '
    else:
        prereqs = \
            'import nuke; ctrl = nuke.toNode(\'' + ctrl_node_name + '\'); '
    if prereqs not in knob_text:
        knob_text += prereqs
    nuke.root().knob('onScriptLoad').setText(knob_text)


def _register_function_with_onScriptLoad(function_name, ctrl_node_name=None):
    """Registers the function with ``nuke.root()['onScriptLoad']``"""

    __register_onScriptLoad_prerequisites(ctrl_node_name)
    knob_text = nuke.root().knob('onScriptLoad').getValue()
    on_script_load = ("exec(ctrl['" + function_name + "'].getText()); " +
                      "print('onScriptLoad executed.')")
    if on_script_load not in knob_text:
        knob_text += on_script_load
    nuke.root().knob('onScriptLoad').setText(knob_text)


def store_functions(functions, ctrl_node_name=None):
    """Add functions onto control node and into onScriptLoad

    If ctrl_node is None, store on nuke.root() instead.
    """

    for function_name, function_text in functions.items():
        _add_function_to_ctrl_node(function_name, function_text, ctrl_node_name)
        _register_function_with_onScriptLoad(function_name, ctrl_node_name)