Visualize pipeline objects in notebook

.github/actions/install_node_dependencies/action.yml

@@ -26,7 +26,7 @@ runs:
       shell: bash
 
     - name: Cache Node.js packages
-      uses: actions/cache@v2
+      uses: actions/cache@v4
       with:
         path: "${{ steps.npm-cache-dir.outputs.dir }}"
         key: "${{ runner.os }}-node-${{ hashFiles(format('{0}/package-lock.json', inputs.package-path)) }}"

.github/workflows/lint.yml

@@ -37,9 +37,6 @@ jobs:
 
           - name: Run security scan
             run: make security-scan
-
-          - name: Verify GraphQL schema is up to date
-            run: make schema-check
-
+
           - name: Run Python formatters and linters
             run: make format-check lint-check          

Makefile

@@ -39,13 +39,6 @@ lint-check:
 	mypy --config-file=package/mypy.ini package/kedro_viz package/features
 	mypy --disable-error-code abstract --config-file=package/mypy.ini package/tests
 
-schema-fix:
-	strawberry export-schema --app-dir=package kedro_viz.api.graphql.schema > src/apollo/schema.graphql
-	graphqlviz src/apollo/schema.graphql | dot -Tpng -o .github/img/schema.graphql.png
-
-schema-check:
-	strawberry export-schema --app-dir=package kedro_viz.api.graphql.schema | diff src/apollo/schema.graphql -
-
 secret-scan:
 	trufflehog --max_depth 1 --exclude_path trufflehog-ignore.txt .
 

RELEASE.md

@@ -9,6 +9,8 @@
 
 ## Major features and improvements
 
+-  Visualize pipeline objects in notebook. (#2241)
+
 ## Bug fixes and other changes
 
 - Add ESM bundle for Kedro-Viz. (#2268)

package/kedro_viz/data_access/managers.py

@@ -52,6 +52,10 @@ class DataAccessManager:
     """Centralised interface for the rest of the application to interact with data repositories."""
 
     def __init__(self):
+        self._initialize_fields()
+
+    def _initialize_fields(self):
+        """Initialize or reset all instance variables."""
         self.catalog = CatalogRepository()
         self.nodes = GraphNodesRepository()
         self.registered_pipelines = RegisteredPipelinesRepository()
@@ -72,6 +76,10 @@ def __init__(self):
         self.tracking_datasets = TrackingDatasetsRepository()
         self.dataset_stats = {}
 
+    def reset_fields(self):
+        """Reset all instance variables."""
+        self._initialize_fields()
+
     def set_db_session(self, db_session_class: sessionmaker):
         """Set db session on repositories that need it."""
         self.runs.set_db_session(db_session_class)

package/kedro_viz/integrations/notebook/__init__.py

@@ -0,0 +1,4 @@
+"""`kedro_viz.integrations.notebook` provides interface to integrate Kedro-Viz with Notebook."""
+
+# alias to ease Notebook visualization import
+from .visualizer import NotebookVisualizer

package/kedro_viz/integrations/notebook/data_loader.py

@@ -0,0 +1,55 @@
+"""`kedro_viz.integrations.notebook.data_loader` provides interface to
+load data from a notebook. It takes care of making sure viz can
+load data from pipelines created in a range of Kedro versions.
+"""
+
+from typing import Dict, Optional, Tuple, Union, cast
+
+from kedro import __version__
+from kedro.framework.session.store import BaseSessionStore
+from kedro.io import DataCatalog
+from kedro.pipeline import Pipeline
+
+from kedro_viz.data_access import data_access_manager
+from kedro_viz.server import populate_data
+
+
+def load_data_for_notebook_users(
+    notebook_pipeline: Union[Pipeline, Dict[str, Pipeline]],
+    notebook_catalog: Optional[DataCatalog],
+) -> Tuple[DataCatalog, Dict[str, Pipeline], BaseSessionStore, Dict]:
+    """Load data from a notebook user's pipeline"""
+    # Create a dummy data catalog with all datasets as memory datasets
+    catalog = DataCatalog() if notebook_catalog is None else notebook_catalog
+    session_store = None
+    stats_dict: Dict = {}
+
+    notebook_user_pipeline = notebook_pipeline
+
+    # create a default pipeline if a dictionary of pipelines are sent
+    if isinstance(notebook_user_pipeline, dict):
+        notebook_user_pipeline = {
+            "__default__": notebook_user_pipeline["__default__"]
+            if "__default__" in notebook_user_pipeline
+            else cast(Pipeline, sum(notebook_user_pipeline.values()))
+        }
+    else:
+        notebook_user_pipeline = {"__default__": notebook_user_pipeline}
+
+    return catalog, notebook_user_pipeline, session_store, stats_dict  # type: ignore[return-value]
+
+
+def load_and_populate_data_for_notebook_users(
+    notebook_pipeline: Union[Pipeline, Dict[str, Pipeline]],
+    notebook_catalog: Optional[DataCatalog],
+):
+    """Loads pipeline data and populates Kedro Viz Repositories for a notebook user"""
+    catalog, pipelines, session_store, stats_dict = load_data_for_notebook_users(
+        notebook_pipeline, notebook_catalog
+    )
+
+    # make each cell independent
+    data_access_manager.reset_fields()
+
+    # Creates data repositories which are used by Kedro Viz Backend APIs
+    populate_data(data_access_manager, catalog, pipelines, session_store, stats_dict)

package/kedro_viz/integrations/notebook/visualizer.py

@@ -0,0 +1,193 @@
+import json
+import logging
+import uuid
+from contextlib import contextmanager
+from typing import Any, Dict, Optional, Union
+
+from IPython.display import HTML, display
+from kedro.io.data_catalog import DataCatalog
+from kedro.pipeline import Pipeline
+
+from kedro_viz.api.rest.responses.pipelines import get_kedro_project_json_data
+from kedro_viz.integrations.notebook.data_loader import (
+    load_and_populate_data_for_notebook_users,
+)
+from kedro_viz.utils import Spinner, merge_dicts
+
+DEFAULT_VIZ_OPTIONS = {
+    "display": {
+        "expandPipelinesBtn": False,
+        "globalNavigation": False,
+        "exportBtn": False,
+        "labelBtn": False,
+        "layerBtn": False,
+        "metadataPanel": False,
+        "miniMap": False,
+        "sidebar": False,
+        "zoomToolbar": False,
+    },
+    "expandAllPipelines": False,
+    "behaviour": {
+        "reFocus": False,
+    },
+    "theme": "dark",
+    "width": "100%",
+    "height": "600px",
+}
+
+DEFAULT_JS_URL = (
+    "https://cdn.jsdelivr.net/gh/kedro-org/kedro-viz@main/esm/kedro-viz.production.mjs"
+)
+
+
+class NotebookVisualizer:
+    """Represent a Kedro-Viz visualization instance in a notebook"""
+
+    def __init__(
+        self,
+        pipeline: Union[Pipeline, Dict[str, Pipeline]],
+        catalog: Optional[DataCatalog] = None,
+        options: Optional[Dict[str, Any]] = None,
+        js_url: Optional[str] = None,
+    ):
+        """
+        Initialize NotebookVisualizer.
+
+        Args:
+            pipeline: Kedro pipeline(s) to visualize.
+            catalog: Kedro data catalog.
+            options: Visualization options.
+            (Ref: https://github.com/kedro-org/kedro-viz/blob/main/README.npm.md#configure-kedro-viz-with-options)
+            js_url: Optional URL for the Kedro-Viz JS bundle.
+
+        Returns:
+            A new ``NotebookVisualizer`` instance.
+        """
+        self.pipeline = pipeline
+        self.catalog = catalog
+        self.options = (
+            DEFAULT_VIZ_OPTIONS
+            if options is None
+            else merge_dicts(DEFAULT_VIZ_OPTIONS, options)
+        )
+        # Force `globalNavigation` to always be False as it
+        # breaks visualizer due to security concerns
+        self.options.setdefault("display", {})["globalNavigation"] = False  # type: ignore
+
+        self.js_url = js_url or DEFAULT_JS_URL
+
+    def _load_viz_data(self) -> Optional[Any]:
+        """Load pipeline and catalog data for visualization."""
+        load_and_populate_data_for_notebook_users(self.pipeline, self.catalog)
+        return get_kedro_project_json_data()
+
+    @staticmethod
+    def generate_html(
+        json_to_visualize: Optional[Any],
+        options: Dict[str, Any] = DEFAULT_VIZ_OPTIONS,
+        js_url: str = DEFAULT_JS_URL,
+    ) -> str:
+        """Generate HTML markup for Kedro-Viz.
+
+        Args:
+            json_to_visualize: Kedro project pipeline data as a json object.
+            options: Visualization options.
+            js_url: Optional URL for the Kedro-Viz JS bundle.
+
+        Returns:
+            The HTML markup template as a string
+        """
+        unique_id = uuid.uuid4().hex[:8]  # To isolate container for each cell execution
+        json_data_str = json.dumps(json_to_visualize)
+        options_str = json.dumps(options)
+
+        html_content = (
+            r"""<!DOCTYPE html>
+        <html lang='en'>
+        <head>
+            <meta charset='UTF-8'>
+            <meta name='viewport' content='width=device-width, initial-scale=1.0'>
+            <title>Kedro-Viz</title>
+        </head>
+        <body>
+            <div id=kedro-viz-"""
+            + unique_id
+            + """ style='height: 600px'></div>
+            <script type="module">
+                import { KedroViz, React, createRoot } from '"""
+            + js_url
+            + """';
+                const viz_container = document.getElementById('kedro-viz-"""
+            + unique_id
+            + """');
+
+                if (createRoot && viz_container) {
+                    const viz_root = createRoot(viz_container);
+                    viz_root.render(
+                        React.createElement(KedroViz, {
+                            data: """
+            + json_data_str
+            + """,
+                            options: """
+            + options_str
+            + """
+                        })
+                    );
+                }
+            </script>
+        </body>
+        </html>"""
+        )
+
+        return html_content
+
+    @staticmethod
+    def _wrap_in_iframe(
+        html_content: str,
+        width: str = str(DEFAULT_VIZ_OPTIONS.get("width", "")),
+        height: str = str(DEFAULT_VIZ_OPTIONS.get("height", "")),
+    ) -> str:
+        """Wrap the HTML content in an iframe.
+
+        Args:
+            html_content: The HTML markup template as a string for visualization
+            width: iframe width
+            height: iframe height
+
+        Returns:
+            A string containing html markup embedded in an iframe
+        """
+        sanitized_content = html_content.replace('"', "&quot;")
+        return f"""<iframe srcdoc="{sanitized_content}"  style="width:{width}; height:{height}; border:none;" sandbox="allow-scripts"></iframe>"""
+
+    @staticmethod
+    @contextmanager
+    def _suppress_logs():
+        logger = logging.getLogger()
+        previous_level = logger.level
+        logger.setLevel(logging.CRITICAL)  # Suppress logs
+        try:
+            yield
+        finally:
+            logger.setLevel(previous_level)  # Restore the original level
+
+    def show(self) -> None:
+        """Display Kedro-Viz in a notebook."""
+        with self._suppress_logs():
+            try:
+                spinner = Spinner("Starting Kedro-Viz...")
+                spinner.start()
+                json_to_visualize = self._load_viz_data()
+                html_content = self.generate_html(
+                    json_to_visualize, self.options, self.js_url
+                )
+                iframe_content = self._wrap_in_iframe(
+                    html_content,
+                    str(self.options.get("width", "")),
+                    str(self.options.get("height", "")),
+                )
+                spinner.stop()
+                display(HTML(iframe_content))
+            except Exception as exc:  # noqa: BLE001
+                spinner.stop()
+                display(HTML(f"<strong>Error: {str(exc)}</strong>"))

package/kedro_viz/utils.py

@@ -1,7 +1,11 @@
 """Transcoding related utility functions."""
 
 import hashlib
-from typing import Tuple
+import sys
+import threading
+import time
+from itertools import cycle
+from typing import Any, Tuple
 
 TRANSCODING_SEPARATOR = "@"
 
@@ -57,3 +61,43 @@ def _strip_transcoding(element: str) -> str:
 def is_dataset_param(dataset_name: str) -> bool:
     """Return whether a dataset is a parameter"""
     return dataset_name.lower().startswith("params:") or dataset_name == "parameters"
+
+
+def merge_dicts(dict_one: dict[str, Any], dict_two: dict[str, Any]) -> dict[str, Any]:
+    """Utility to merge two dictionaries"""
+    import copy
+
+    merged = copy.deepcopy(dict_one)
+
+    for key, value in dict_two.items():
+        if isinstance(value, dict) and key in merged:
+            merged[key] = merge_dicts(merged[key], value)
+        else:
+            merged[key] = value
+    return merged
+
+
+class Spinner:
+    """Represent a simple spinner instance"""
+
+    def __init__(self, message: str = "Processing"):
+        self.spinner = cycle(["-", "\\", "|", "/"])
+        self.message = message
+        self.stop_running = False
+
+    def start(self):
+        def run_spinner():
+            while not self.stop_running:
+                sys.stdout.write(f"\r{self.message} {next(self.spinner)} ")
+                sys.stdout.flush()
+                time.sleep(0.1)
+            sys.stdout.write(
+                "\r" + " " * (len(self.message) + 2) + "\r"
+            )  # Clear the line
+
+        self._spinner_thread = threading.Thread(target=run_spinner, daemon=True)
+        self._spinner_thread.start()
+
+    def stop(self):
+        self.stop_running = True
+        self._spinner_thread.join()