Refactor CLI with lazy subcommands and deferring imports (#1920)

* try delaying import + lazy subcommands

Signed-off-by: Ankita Katiyar <[email protected]>

* lazy group initial draft

* fix lint

* add tests and fix lint

* add default command and tests

* add default command and tests

* address PR comments

---------

Signed-off-by: Ankita Katiyar <[email protected]>
Co-authored-by: ravi-kumar-pilla <[email protected]>

package/kedro_viz/launchers/cli/__init__.py

@@ -0,0 +1 @@
+"""`kedro_viz.launchers.cli` launches the viz server as a CLI app."""

package/kedro_viz/launchers/cli/build.py

@@ -0,0 +1,24 @@
+"""`kedro_viz.launchers.cli.build` provides a cli command to build
+a Kedro-Viz instance"""
+# pylint: disable=import-outside-toplevel
+import click
+
+from kedro_viz.launchers.cli.main import viz
+
+
+@viz.command(context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "--include-hooks",
+    is_flag=True,
+    help="A flag to include all registered hooks in your Kedro Project",
+)
+@click.option(
+    "--include-previews",
+    is_flag=True,
+    help="A flag to include preview for all the datasets",
+)
+def build(include_hooks, include_previews):
+    """Create build directory of local Kedro Viz instance with Kedro project data"""
+    from kedro_viz.launchers.cli.utils import create_shareableviz_process
+
+    create_shareableviz_process("local", include_previews, include_hooks=include_hooks)

package/kedro_viz/launchers/cli/deploy.py

@@ -0,0 +1,69 @@
+"""`kedro_viz.launchers.cli.deploy` provides a cli command to deploy
+a Kedro-Viz instance on cloud platforms"""
+# pylint: disable=import-outside-toplevel
+import click
+
+from kedro_viz.constants import SHAREABLEVIZ_SUPPORTED_PLATFORMS
+from kedro_viz.launchers.cli.main import viz
+
+
+@viz.command(context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "--platform",
+    type=str,
+    required=True,
+    help=f"Supported Cloud Platforms like {*SHAREABLEVIZ_SUPPORTED_PLATFORMS,} to host Kedro Viz",
+)
+@click.option(
+    "--endpoint",
+    type=str,
+    required=True,
+    help="Static Website hosted endpoint."
+    "(eg., For AWS - http://<bucket_name>.s3-website.<region_name>.amazonaws.com/)",
+)
+@click.option(
+    "--bucket-name",
+    type=str,
+    required=True,
+    help="Bucket name where Kedro Viz will be hosted",
+)
+@click.option(
+    "--include-hooks",
+    is_flag=True,
+    help="A flag to include all registered hooks in your Kedro Project",
+)
+@click.option(
+    "--include-previews",
+    is_flag=True,
+    help="A flag to include preview for all the datasets",
+)
+def deploy(platform, endpoint, bucket_name, include_hooks, include_previews):
+    """Deploy and host Kedro Viz on provided platform"""
+    from kedro_viz.launchers.cli.utils import (
+        create_shareableviz_process,
+        display_cli_message,
+    )
+
+    if not platform or platform.lower() not in SHAREABLEVIZ_SUPPORTED_PLATFORMS:
+        display_cli_message(
+            "ERROR: Invalid platform specified. Kedro-Viz supports \n"
+            f"the following platforms - {*SHAREABLEVIZ_SUPPORTED_PLATFORMS,}",
+            "red",
+        )
+        return
+
+    if not endpoint:
+        display_cli_message(
+            "ERROR: Invalid endpoint specified. If you are looking for platform \n"
+            "agnostic shareable viz solution, please use the `kedro viz build` command",
+            "red",
+        )
+        return
+
+    create_shareableviz_process(
+        platform,
+        include_previews,
+        endpoint,
+        bucket_name,
+        include_hooks,
+    )

package/kedro_viz/launchers/cli/lazy_default_group.py

@@ -0,0 +1,76 @@
+"""`kedro_viz.launchers.cli.lazy_default_group` provides a custom mutli-command
+subclass for a lazy subcommand loader"""
+
+# pylint: disable=import-outside-toplevel
+from typing import Any, Union
+
+import click
+
+
+class LazyDefaultGroup(click.Group):
+    """A click Group that supports lazy loading of subcommands and a default command"""
+
+    def __init__(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ):
+        if not kwargs.get("ignore_unknown_options", True):
+            raise ValueError("Default group accepts unknown options")
+        self.ignore_unknown_options = True
+
+        # lazy_subcommands is a map of the form:
+        #
+        #   {command-name} -> {module-name}.{command-object-name}
+        #
+        self.lazy_subcommands = kwargs.pop("lazy_subcommands", {})
+
+        self.default_cmd_name = kwargs.pop("default", None)
+        self.default_if_no_args = kwargs.pop("default_if_no_args", False)
+
+        super().__init__(*args, **kwargs)
+
+    def list_commands(self, ctx: click.Context) -> list[str]:
+        return sorted(self.lazy_subcommands.keys())
+
+    def get_command(  # type: ignore[override]
+        self, ctx: click.Context, cmd_name: str
+    ) -> Union[click.BaseCommand, click.Command, None]:
+        if cmd_name in self.lazy_subcommands:
+            return self._lazy_load(cmd_name)
+        return super().get_command(ctx, cmd_name)
+
+    def _lazy_load(self, cmd_name: str) -> click.BaseCommand:
+        from importlib import import_module
+
+        # lazily loading a command, first get the module name and attribute name
+        import_path = self.lazy_subcommands[cmd_name]
+        modname, cmd_object_name = import_path.rsplit(".", 1)
+
+        # do the import
+        mod = import_module(modname)
+
+        # get the Command object from that module
+        cmd_object = getattr(mod, cmd_object_name)
+
+        return cmd_object
+
+    def parse_args(self, ctx, args):
+        # If no args are provided and default_command_name is specified,
+        # use the default command
+        if not args and self.default_if_no_args:
+            args.insert(0, self.default_cmd_name)
+        return super().parse_args(ctx, args)
+
+    def resolve_command(self, ctx: click.Context, args):
+        # Attempt to resolve the command using the parent class method
+        try:
+            cmd_name, cmd, args = super().resolve_command(ctx, args)
+            return cmd_name, cmd, args
+        except click.UsageError as exc:
+            if self.default_cmd_name and not ctx.invoked_subcommand:
+                # No command found, use the default command
+                default_cmd = self.get_command(ctx, self.default_cmd_name)
+                if default_cmd:
+                    return default_cmd.name, default_cmd, args
+            raise exc

package/kedro_viz/launchers/cli/main.py

@@ -0,0 +1,26 @@
+"""`kedro_viz.launchers.cli.main` is an entry point for Kedro-Viz cli commands."""
+
+import click
+
+from kedro_viz.launchers.cli.lazy_default_group import LazyDefaultGroup
+
+
+@click.group(name="Kedro-Viz")
+def viz_cli():  # pylint: disable=missing-function-docstring
+    pass
+
+
+@viz_cli.group(
+    name="viz",
+    cls=LazyDefaultGroup,
+    lazy_subcommands={
+        "run": "kedro_viz.launchers.cli.run.run",
+        "deploy": "kedro_viz.launchers.cli.deploy.deploy",
+        "build": "kedro_viz.launchers.cli.build.build",
+    },
+    default="run",
+    default_if_no_args=True,
+)
+@click.pass_context
+def viz(ctx):  # pylint: disable=unused-argument
+    """Visualise a Kedro pipeline using Kedro viz."""

package/kedro_viz/launchers/cli/run.py

@@ -0,0 +1,196 @@
+"""`kedro_viz.launchers.cli.run` provides a cli command to run
+a Kedro-Viz instance"""
+
+from typing import Dict
+
+import click
+from kedro.framework.cli.project import PARAMS_ARG_HELP
+from kedro.framework.cli.utils import _split_params
+
+from kedro_viz.constants import DEFAULT_HOST, DEFAULT_PORT
+from kedro_viz.launchers.cli.main import viz
+
+_VIZ_PROCESSES: Dict[str, int] = {}
+
+
+@viz.command(context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "--host",
+    default=DEFAULT_HOST,
+    help="Host that viz will listen to. Defaults to localhost.",
+)
+@click.option(
+    "--port",
+    default=DEFAULT_PORT,
+    type=int,
+    help="TCP port that viz will listen to. Defaults to 4141.",
+)
+@click.option(
+    "--browser/--no-browser",
+    default=True,
+    help="Whether to open viz interface in the default browser or not. "
+    "Browser will only be opened if host is localhost. Defaults to True.",
+)
+@click.option(
+    "--load-file",
+    default=None,
+    help="Path to load Kedro-Viz data from a directory",
+)
+@click.option(
+    "--save-file",
+    default=None,
+    type=click.Path(dir_okay=False, writable=True),
+    help="Path to save Kedro-Viz data to a directory",
+)
+@click.option(
+    "--pipeline",
+    "-p",
+    type=str,
+    default=None,
+    help="Name of the registered pipeline to visualise. "
+    "If not set, the default pipeline is visualised",
+)
+@click.option(
+    "--env",
+    "-e",
+    type=str,
+    default=None,
+    multiple=False,
+    envvar="KEDRO_ENV",
+    help="Kedro configuration environment. If not specified, "
+    "catalog config in `local` will be used",
+)
+@click.option(
+    "--autoreload",
+    "-a",
+    is_flag=True,
+    help="Autoreload viz server when a Python or YAML file change in the Kedro project",
+)
+@click.option(
+    "--include-hooks",
+    is_flag=True,
+    help="A flag to include all registered hooks in your Kedro Project",
+)
+@click.option(
+    "--params",
+    type=click.UNPROCESSED,
+    default="",
+    help=PARAMS_ARG_HELP,
+    callback=_split_params,
+)
+# pylint: disable=import-outside-toplevel, too-many-locals
+def run(
+    host,
+    port,
+    browser,
+    load_file,
+    save_file,
+    pipeline,
+    env,
+    autoreload,
+    include_hooks,
+    params,
+):
+    """Launch local Kedro Viz instance"""
+    # Deferring Imports
+    import multiprocessing
+    import traceback
+    from pathlib import Path
+
+    from kedro.framework.cli.utils import KedroCliError
+    from kedro.framework.project import PACKAGE_NAME
+    from packaging.version import parse
+
+    from kedro_viz import __version__
+    from kedro_viz.integrations.pypi import (
+        get_latest_version,
+        is_running_outdated_version,
+    )
+    from kedro_viz.launchers.cli.utils import display_cli_message
+    from kedro_viz.launchers.utils import (
+        _PYPROJECT,
+        _check_viz_up,
+        _find_kedro_project,
+        _start_browser,
+        _wait_for,
+    )
+    from kedro_viz.server import run_server
+
+    kedro_project_path = _find_kedro_project(Path.cwd())
+
+    if kedro_project_path is None:
+        display_cli_message(
+            "ERROR: Failed to start Kedro-Viz : "
+            "Could not find the project configuration "
+            f"file '{_PYPROJECT}' at '{Path.cwd()}'. ",
+            "red",
+        )
+        return
+
+    installed_version = parse(__version__)
+    latest_version = get_latest_version()
+    if is_running_outdated_version(installed_version, latest_version):
+        display_cli_message(
+            "WARNING: You are using an old version of Kedro Viz. "
+            f"You are using version {installed_version}; "
+            f"however, version {latest_version} is now available.\n"
+            "You should consider upgrading via the `pip install -U kedro-viz` command.\n"
+            "You can view the complete changelog at "
+            "https://github.com/kedro-org/kedro-viz/releases.",
+            "yellow",
+        )
+    try:
+        if port in _VIZ_PROCESSES and _VIZ_PROCESSES[port].is_alive():
+            _VIZ_PROCESSES[port].terminate()
+
+        run_server_kwargs = {
+            "host": host,
+            "port": port,
+            "load_file": load_file,
+            "save_file": save_file,
+            "pipeline_name": pipeline,
+            "env": env,
+            "project_path": kedro_project_path,
+            "autoreload": autoreload,
+            "include_hooks": include_hooks,
+            "package_name": PACKAGE_NAME,
+            "extra_params": params,
+        }
+        if autoreload:
+            from watchgod import RegExpWatcher, run_process
+
+            run_process_kwargs = {
+                "path": kedro_project_path,
+                "target": run_server,
+                "kwargs": run_server_kwargs,
+                "watcher_cls": RegExpWatcher,
+                "watcher_kwargs": {"re_files": r"^.*(\.yml|\.yaml|\.py|\.json)$"},
+            }
+            viz_process = multiprocessing.Process(
+                target=run_process, daemon=False, kwargs={**run_process_kwargs}
+            )
+        else:
+            viz_process = multiprocessing.Process(
+                target=run_server, daemon=False, kwargs={**run_server_kwargs}
+            )
+
+        display_cli_message("Starting Kedro Viz ...", "green")
+
+        viz_process.start()
+
+        _VIZ_PROCESSES[port] = viz_process
+
+        _wait_for(func=_check_viz_up, host=host, port=port)
+
+        display_cli_message(
+            "Kedro Viz started successfully. \n\n"
+            f"\u2728 Kedro Viz is running at \n http://{host}:{port}/",
+            "green",
+        )
+
+        if browser:
+            _start_browser(host, port)
+
+    except Exception as ex:  # pragma: no cover
+        traceback.print_exc()
+        raise KedroCliError(str(ex)) from ex