Lightweight Kedro Viz Experimentation using AST

package/kedro_viz/integrations/kedro/data_loader.py

@@ -11,14 +11,14 @@
 from typing import Any, Dict, Optional, Tuple
 
 from kedro import __version__
-from kedro.framework.project import configure_project, pipelines
-from kedro.framework.session import KedroSession
+from kedro.config.omegaconf_config import OmegaConfigLoader
+from kedro.framework.context.context import KedroContext
 from kedro.framework.session.store import BaseSessionStore
-from kedro.framework.startup import bootstrap_project
 from kedro.io import DataCatalog
 from kedro.pipeline import Pipeline
 
 from kedro_viz.constants import VIZ_METADATA_ARGS
+from kedro_viz.integrations.kedro.lite_parser import parse_project
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +75,7 @@ def load_data(
     include_hooks: bool = False,
     package_name: Optional[str] = None,
     extra_params: Optional[Dict[str, Any]] = None,
+    is_lite: bool = False,
 ) -> Tuple[DataCatalog, Dict[str, Pipeline], BaseSessionStore, Dict]:
     """Load data from a Kedro project.
     Args:
@@ -91,30 +92,54 @@ def load_data(
         A tuple containing the data catalog and the pipeline dictionary
         and the session store.
     """
-    if package_name:
-        configure_project(package_name)
-    else:
-        # bootstrap project when viz is run in dev mode
-        bootstrap_project(project_path)
-
-    with KedroSession.create(
-        project_path=project_path,
-        env=env,
-        save_on_close=False,
-        extra_params=extra_params,
-    ) as session:
-        # check for --include-hooks option
-        if not include_hooks:
-            session._hook_manager = _VizNullPluginManager()  # type: ignore
-
-        context = session.load_context()
-        session_store = session._store
-        catalog = context.catalog
-
-        # Pipelines is a lazy dict-like object, so we force it to populate here
-        # in case user doesn't have an active session down the line when it's first accessed.
-        # Useful for users who have `get_current_session` in their `register_pipelines()`.
-        pipelines_dict = dict(pipelines)
-        stats_dict = _get_dataset_stats(project_path)
+    if is_lite:
+        # [TODO: Confirm on the context creation]
+        context = KedroContext(
+            package_name="{{ cookiecutter.python_package }}",
+            project_path=project_path,
+            config_loader=OmegaConfigLoader(conf_source=str(project_path)),
+            hook_manager=_VizNullPluginManager(),
+            env=env,
+        )
+
+        # [TODO: Confirm on the session store creation]
+        session_store = None
 
-    return catalog, pipelines_dict, session_store, stats_dict
+        # [TODO: Confirm on the DataCatalog creation]
+        catalog = DataCatalog()
+
+        stats_dict = _get_dataset_stats(project_path)
+        pipelines_dict = dict(parse_project(project_path))
+        return catalog, pipelines_dict, session_store, stats_dict
+    else:
+        from kedro.framework.project import configure_project, pipelines
+        from kedro.framework.session import KedroSession
+        from kedro.framework.startup import bootstrap_project
+
+        if package_name:
+            configure_project(package_name)
+        else:
+            # bootstrap project when viz is run in dev mode
+            bootstrap_project(project_path)
+
+        with KedroSession.create(
+            project_path=project_path,
+            env=env,
+            save_on_close=False,
+            extra_params=extra_params,
+        ) as session:
+            # check for --include-hooks option
+            if not include_hooks:
+                session._hook_manager = _VizNullPluginManager()  # type: ignore
+
+            context = session.load_context()
+            session_store = session._store
+            catalog = context.catalog
+
+            # Pipelines is a lazy dict-like object, so we force it to populate here
+            # in case user doesn't have an active session down the line when it's first accessed.
+            # Useful for users who have `get_current_session` in their `register_pipelines()`.
+            pipelines_dict = dict(pipelines)
+            stats_dict = _get_dataset_stats(project_path)
+
+        return catalog, pipelines_dict, session_store, stats_dict

package/kedro_viz/integrations/kedro/lite_parser.py

@@ -0,0 +1,242 @@
+import ast
+import logging
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, Iterable, List
+
+from kedro.pipeline.modular_pipeline import pipeline as ModularPipeline
+from kedro.pipeline.pipeline import Node, Pipeline
+
+logger = logging.getLogger(__name__)
+
+
+class KedroPipelineLocator(ast.NodeVisitor):
+    """
+    Represents a pipeline that is located when parsing
+    the Kedro project's `create_pipeline` function
+
+    """
+
+    def __init__(self):
+        self.pipeline = None
+
+    def visit_FunctionDef(self, node):
+        try:
+            if node.name == "create_pipeline":
+                # Explore the located pipeline for nodes
+                # and other keyword args
+                kedro_pipeline_explorer = KedroPipelineExplorer()
+                kedro_pipeline_explorer.visit(node)
+                try:
+                    # modular pipeline
+                    if kedro_pipeline_explorer.namespace:
+                        self.pipeline = ModularPipeline(
+                            pipe=kedro_pipeline_explorer.nodes,
+                            inputs=kedro_pipeline_explorer.inputs,
+                            outputs=kedro_pipeline_explorer.outputs,
+                            parameters=kedro_pipeline_explorer.parameters,
+                            tags=kedro_pipeline_explorer.tags,
+                            namespace=kedro_pipeline_explorer.namespace,
+                        )
+                    else:
+                        # kedro pipeline
+                        self.pipeline = Pipeline(
+                            nodes=kedro_pipeline_explorer.nodes,
+                            tags=kedro_pipeline_explorer.tags,
+                        )
+                except Exception as exc:
+                    # [TODO: Error with modular pipeline, try creating regular pipeline]
+                    logger.error(exc)
+                    self.pipeline = Pipeline(
+                        nodes=kedro_pipeline_explorer.nodes,
+                        tags=kedro_pipeline_explorer.tags,
+                    )
+
+            self.generic_visit(node)
+
+        except Exception as exc:
+            # [TODO: Error with parsing the file,
+            # dump the visiting node for debugging]
+            logger.error(exc)
+            logger.info(ast.dump(node, indent=2))
+
+
+class KedroPipelineExplorer(ast.NodeVisitor):
+    # [TODO: Current explorer only serves for 1 pipeline() function within a create_pipeline def]
+    def __init__(self):
+        # keeping these here for future use-case
+        # when dealing with multiple `pipeline()` functions
+        # within a create_pipeline def
+        self.nodes: List[Node] = []
+        self.inputs = None
+        self.outputs = None
+        self.namespace = None
+        self.parameters = None
+        self.tags = None
+
+    def visit_Call(self, node):
+        if isinstance(node.func, ast.Name) and node.func.id == "pipeline":
+            # for a modular pipeline
+            # [TODO: pipe to be explored later]
+            # pipe: Iterable[Node | Pipeline] | Pipeline
+
+            pipeline_inputs: str | set[str] | dict[str, str] | None = None
+            pipeline_outputs: str | set[str] | dict[str, str] | None = None
+            pipeline_namespace: str | None = None
+            pipeline_parameters: str | set[str] | dict[str, str] | None = None
+            pipeline_tags: str | Iterable[str] | None = None
+
+            for keyword in node.keywords:
+                if keyword.arg == "namespace":
+                    pipeline_namespace = parse_value(keyword.value)
+                elif keyword.arg == "inputs":
+                    pipeline_inputs = parse_value(keyword.value)
+                elif keyword.arg == "outputs":
+                    pipeline_outputs = parse_value(keyword.value)
+                elif keyword.arg == "parameters":
+                    pipeline_parameters = parse_value(keyword.value)
+                elif keyword.arg == "tags":
+                    pipeline_tags = parse_value(keyword.value)
+
+            # exploring nodes
+            for arg in node.args:
+                if isinstance(arg, ast.List):
+                    for elt in arg.elts:
+                        if (
+                            isinstance(elt, ast.Call)
+                            and isinstance(elt.func, ast.Name)
+                            and elt.func.id == "node"
+                        ):
+                            node_func = None
+                            node_inputs: str | list[str] | dict[str, str] | None = None
+                            node_outputs: str | list[str] | dict[str, str] | None = None
+                            node_name: str | None = None
+                            node_tags: str | Iterable[str] | None = None
+                            node_confirms: str | list[str] | None = None
+                            node_namespace: str | None = None
+
+                            for keyword in elt.keywords:
+                                # [TODO: func is WIP. Need to create a Callable]
+                                if keyword.arg == "func":
+                                    if isinstance(keyword.value, ast.Name):
+                                        func_name = keyword.value.id
+                                        exec(
+                                            f"def {func_name}(*args, **kwargs): pass",
+                                            globals(),
+                                        )
+                                        node_func = globals()[func_name]
+                                    else:
+                                        node_func = lambda *args, **kwargs: None
+                                elif keyword.arg == "inputs":
+                                    node_inputs = parse_value(keyword.value)
+                                elif keyword.arg == "outputs":
+                                    node_outputs = parse_value(keyword.value)
+                                elif keyword.arg == "name":
+                                    node_name = parse_value(keyword.value)
+                                elif keyword.arg == "tags":
+                                    node_tags = parse_value(keyword.value)
+                                elif keyword.arg == "confirms":
+                                    node_confirms = parse_value(keyword.value)
+                                elif keyword.arg == "namespace":
+                                    node_namespace = parse_value(keyword.value)
+
+                            # Create Node
+                            kedro_node = Node(
+                                func=node_func,
+                                inputs=node_inputs,
+                                outputs=node_outputs,
+                                name=node_name,
+                                tags=node_tags,
+                                confirms=node_confirms,
+                                namespace=node_namespace,
+                            )
+
+                            self.nodes.append(kedro_node)
+
+            # These will be used for modular pipeline creation
+            self.inputs = pipeline_inputs
+            self.outputs = pipeline_outputs
+            self.namespace = pipeline_namespace
+            self.parameters = pipeline_parameters
+            self.tags = pipeline_tags
+
+        self.generic_visit(node)
+
+
+# Helper functions
+def parse_value(keyword_value):
+    """Helper to parse values assigned to node/pipeline properties"""
+    if isinstance(keyword_value, ast.Constant):
+        if not keyword_value.value:
+            return None
+        return str(keyword_value.value)
+    elif isinstance(keyword_value, (ast.List, ast.Set)):
+        return [parse_value(elt) for elt in keyword_value.elts]
+    elif isinstance(keyword_value, ast.Dict):
+        return {
+            parse_value(k): parse_value(v)
+            for k, v in zip(keyword_value.keys, keyword_value.values)
+        }
+    elif isinstance(keyword_value, ast.ListComp):
+        # [TODO: For list comprehensions, complex case handling]
+        # [Example can be found under demo_project/pipelines/modelling]
+        return f"ListComp({ast.dump(keyword_value)})"
+    elif isinstance(keyword_value, ast.DictComp):
+        # [TODO: For dict comprehensions, complex case handling]
+        # [Example can be found under demo_project/pipelines/modelling]
+        return f"DictComp({ast.dump(keyword_value)})"
+    elif isinstance(keyword_value, ast.FormattedValue):
+        # [TODO: For formatted strings i.e., single formatted fields,
+        # complex case handling]
+        # [Example can be found under demo_project/pipelines/modelling]
+        return f"FormattedValue({ast.dump(keyword_value)})"
+    elif isinstance(keyword_value, ast.JoinedStr):
+        # [TODO: For joined strings i.e., multiple formatted fields,
+        # complex case handling]
+        # [Example can be found under demo_project/pipelines/modelling]
+        return f"JoinedStr({ast.dump(keyword_value)})"
+    elif isinstance(keyword_value, ast.Name):
+        # [TODO: For variable references, complex case handling]
+        # [Example can be found under demo_project/pipelines/modelling]
+        return f"Variable({ast.dump(keyword_value)})"
+    else:
+        # [TODO: For any other complex case handling]
+        return f"Unsupported({ast.dump(keyword_value)})"
+
+
+# [WIP: Naive parsing and exploring pipelines. Not sure of any better way for now]
+def parse_project(project_path: Path) -> Dict[str, Pipeline]:
+    # Result
+    pipelines: Dict[str, Pipeline] = defaultdict(dict)
+
+    # Loop through all the .py files in the kedro project
+    # and start locating create_pipeline
+    for filepath in project_path.rglob("*.py"):
+        with open(filepath, "r") as file:
+            file_content = file.read()
+
+        # parse file content using ast
+        parsed_content_ast_node = ast.parse(file_content)
+
+        # extract pipeline name from file path
+        pipeline_name = filepath.relative_to(project_path).parent.name
+
+        # Locate pipelines (tested for only 1 create_pipeline per pipeline file)
+        # [TODO: confirm with Kedro team if more than 1 create_pipeline existence]
+        kedro_pipeline_locator = KedroPipelineLocator()
+        kedro_pipeline_locator.visit(parsed_content_ast_node)
+        located_pipeline = kedro_pipeline_locator.pipeline
+
+        # add to the result if a pipeline is located
+        if located_pipeline:
+            pipelines[pipeline_name] = located_pipeline
+
+    # foolproof to have atleast 1 pipeline
+    # so the UI won't break
+    if len(pipelines.keys()):
+        # creating a default pipeline
+        pipelines["__default__"] = sum(pipelines.values())
+    else:
+        pipelines["__default__"] = Pipeline(nodes=[])
+
+    return pipelines

package/kedro_viz/launchers/cli.py

@@ -116,6 +116,11 @@ def viz(ctx):  # pylint: disable=unused-argument
     help=PARAMS_ARG_HELP,
     callback=_split_params,
 )
+@click.option(
+    "--lite",
+    is_flag=True,
+    help="A flag to load an experimental light-weight Kedro Viz",
+)
 # pylint: disable=import-outside-toplevel, too-many-locals
 def run(
     host,
@@ -128,6 +133,7 @@ def run(
     autoreload,
     include_hooks,
     params,
+    lite,
 ):
     """Launch local Kedro Viz instance"""
     from kedro_viz.server import run_server
@@ -171,6 +177,7 @@ def run(
             "include_hooks": include_hooks,
             "package_name": PACKAGE_NAME,
             "extra_params": params,
+            "is_lite": lite,
         }
         if autoreload:
             run_process_kwargs = {