kedro-debug-hook.py

""" Kedro Debug Hook module """
from kedro.framework.hooks import hook_impl

class debug_hook:
    """ Kedro Debug Hook

    Opens a debugger at any hook-able point of your kedro projects lifecycle.
    debug_hook is applied by adding it to the pipeline and setting the desired
    debug points to true.

    Examples:

        >>> hooks = [debug_hook(should_debug_all=True)]
        >>> hooks = [debug_hook(should_debug_before_pipeline_run=True)]

    Args:
        should_debug_all (bool): overrides all points Defaults to False
        should_debug_before_pipeline_run (bool): opens a debugger
            before_pipeline_run if True Defaults to False
        should_debug_after_pipeline_run (bool): opens a debugger
            after_pipeline_run if True Defaults to False
        should_debug_on_pipeline_error (bool): opens a debugger
            on_pipeline_error if True Defaults to False
        should_debug_before_node_run (bool): opens a debugger
            before_node_run if True Defaults to False
        should_debug_after_node_run (bool): opens a debugger
            after_node_run if True Defaults to False
        should_debug_on_node_error (bool): opens a debugger
            on_node_error if True Defaults to False
        should_debug_after_catalog_created (bool): opens a debugger
            after_catalog_created if True Defaults to False

    """
    def __init__(
        self,
        should_debug_all=False,
        should_debug_before_pipeline_run=False,
        should_debug_after_pipeline_run=False,
        should_debug_on_pipeline_error=False,
        should_debug_before_node_run=False,
        should_debug_after_node_run=False,
        should_debug_on_node_error=False,
        should_debug_after_catalog_created=False,
    ):
        self.should_debug_before_pipeline_run = (
            should_debug_before_pipeline_run or should_debug_all
        )
        self.should_debug_after_pipeline_run = (
            should_debug_after_pipeline_run or should_debug_all
        )
        self.should_debug_on_pipeline_error = (
            should_debug_on_pipeline_error or should_debug_all
        )
        self.should_debug_before_node_run = (
            should_debug_before_node_run or should_debug_all
        )
        self.should_debug_after_node_run = (
            should_debug_after_node_run or should_debug_all
        )
        self.should_debug_on_node_error = should_debug_on_node_error or should_debug_all
        self.should_debug_after_catalog_created = (
            should_debug_after_catalog_created or should_debug_all
        )

    @hook_impl
    def before_pipeline_run(self, run_params, pipeline, catalog):
        "pops into a debugger before pipeline run"
        if self.should_debug_before_pipeline_run:
            breakpoint()

    @hook_impl
    def after_pipeline_run(self, run_params, pipeline, catalog):
        "pops into a debugger after pipeline run"
        if self.should_debug_after_pipeline_run:
            breakpoint()

    @hook_impl
    def on_pipeline_error(self, error, run_params, pipeline, catalog):
        "pops into a debugger on pipeline error"
        if self.should_debug_on_pipeline_error:
            breakpoint()

    @hook_impl
    def after_catalog_created(
        self,
        catalog,
        conf_catalog,
        conf_creds,
        feed_dict,
        save_version,
        load_versions,
        run_id,
    ):
        "pops into a debugger after catalog created"
        if self.should_debug_after_catalog_created:
            breakpoint()

    @hook_impl
    def before_node_run(self, node, catalog, inputs, is_async, run_id):
        "pops into a debugger before node run"
        if self.should_debug_before_node_run:
            breakpoint()

    @hook_impl
    def after_node_run(self, node, catalog, inputs, outputs, is_async, run_id):
        "pops into a debugger after node run"
        if self.should_debug_after_node_run:
            breakpoint()

    @hook_impl
    def on_node_error(self, error, node, catalog, inputs, is_async, run_id):
        "pops into a debugger on node error"
        if self.should_debug_on_node_error:
            breakpoint()