revng-revng/python/revng/pypeline/main.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#
from __future__ import annotations

import importlib.util
import os
import sys
from pathlib import Path
from typing import Sequence

import click
import psutil
from click.shell_completion import get_completion_class

from . import initialize_pypeline
from .cli.pipeline import pipeline
from .cli.project import project
from .cli.utils import EagerParsedPath, PypeGroup
from .utils.logger import pypeline_logger


def import_pipebox(module_path: str, is_autocomplete: bool) -> object:
    """Import a module from a path. This is used to import the pipebox file.
    Args:
        env (dict[str, str]): The environment variables.
        module_path (str): The path to the module.
        is_autocomplete (bool): If True, do not raise an error if the pipebox
            is not found, since when autocompleting the pipebox path might be
            missing or incomplete.
    """
    # Absolute path to the module
    module_abspath = Path(module_path).resolve()
    if not module_abspath.exists():
        # This is a small trick to allow generating the module auto-complete
        # without having a pipebox file
        if is_autocomplete:
            return object()
        pypeline_logger.log(
            f'Pipebox file "{module_abspath}" does not exist. Either set it using the '
            "PIPEBOX env var, or pass the --pipebox option."
        )
        sys.exit(1)
    # We guess that the module name is the file name without the extension
    module_name: str = module_abspath.stem
    # Dynamic import of the pipebox module
    spec = importlib.util.spec_from_file_location(module_name, str(module_abspath))
    if spec is None:
        if is_autocomplete:
            return object()
        pypeline_logger.log(f'Could not load module "{module_name}" from "{module_abspath}"')
        sys.exit(1)
    module = importlib.util.module_from_spec(spec)
    if spec.loader is None:
        if is_autocomplete:
            return object()
        pypeline_logger.log(f'Could not load module "{module_name}" from "{module_abspath}"')
        sys.exit(1)
    # Execute the module to load it
    spec.loader.exec_module(module)
    # Initialize the pypeline as we just imported the pipebox
    initialize_pypeline()
    # If auto-completing also initialize the pipebox with empty arguments here
    # instead of in the body of pype
    if is_autocomplete:
        call_pipebox_initialize(module, module_path, [])
    return module


def call_pipebox_initialize(pipebox, pipebox_path: str, args: list[str]):
    # Get its initialize function
    pipebox_initialize = getattr(pipebox, "initialize", None)
    if pipebox_initialize is None:
        pypeline_logger.log(
            f'Pipebox file "{pipebox_path}" does not have an "initialize" function.'
            " This is required to setup the pypeline."
        )
        sys.exit(1)

    # Call the initialize
    pipebox_initialize(args)


def get_current_root_name(ctx: click.Context) -> str:
    root = ctx.find_root()
    command = root.command
    assert command.name is not None, "Command name should not be None"
    return command.name


def detect_autocomplete(ctx: click.Context) -> bool:
    """Detect if we are in auto-complete mode."""
    return (
        "_{}_COMPLETE".format(get_current_root_name(ctx).upper()) in os.environ
        or "autocomplete" in sys.argv
    )


@click.group(cls=PypeGroup)
@click.option(
    "-C",
    "--directory",
    # This impacts `--pipebox` as well, so we need to change the cwd before
    # parsing any other argument
    type=EagerParsedPath(
        name="directory",
        file_okay=False,
        dir_okay=True,
        resolve_path=True,
        parser=lambda path, _ctx: os.chdir(path) if path else None,
    ),
    help=(
        'Change the current working directory, equivalent to running "cd" before executing any'
        " command."
    ),
    expose_value=False,
)
@click.option(
    "--pipebox",
    # The pypebox needs to be imported before the arguments of the subcommands are parsed
    # so we use an eager option to import it as soon as possible
    type=EagerParsedPath(
        name="pipebox",
        # During auto-completion we don't want to fail if the file does not exist
        exists=False,
        parser=lambda path, ctx: import_pipebox(path, detect_autocomplete(ctx)),
    ),
    help=(
        'Path to the pipebox file. Defaults to the "PIPEBOX" environment '
        'variable, then "pipebox.py".'
    ),
    default=os.environ.get("PIPEBOX", "pipebox.py"),
    show_default=True,
    expose_value=False,
)
@click.option(
    "--verbose",
    is_flag=True,
    help="Enable debug logging for the pypeline related code.",
)
@click.pass_context
def pype(ctx, verbose: bool) -> None:

    # Enable debug logging for pypeline if requested
    if verbose:
        pypeline_logger.debug = True

    # Avoid initializing the pipebox if we are in auto-complete mode
    if detect_autocomplete(ctx):
        return

    # Get the already loaded pipebox module from the context
    pipebox = ctx.obj["pipebox"]
    # Initialize the pipebox
    call_pipebox_initialize(pipebox, ctx.obj["pipebox_path"], ctx.obj["pipebox_args"])


pype.add_command(pipeline)
pype.add_command(project)


def detect_shell() -> str:
    """Detect the current shell."""
    try:
        # Get the parent process ID (PPID) of the current Python script
        ppid = os.getppid()
        # Get the process object from the PPID
        parent_process = psutil.Process(ppid)
        # The name of the executable of the parent process is our shell
        return parent_process.name()
    except (psutil.NoSuchProcess, psutil.AccessDenied):
        # Handle cases where the parent process might not exist or is inaccessible
        return Path(os.environ.get("SHELL", "bash")).name


@pype.command()
@click.option(
    "--shell",
    type=click.Choice(["bash", "zsh", "fish"]),
    default=detect_shell(),
    help="Shell type",
)
@click.pass_context
def autocomplete(ctx, shell):
    """
    Generate shell completion script for the CLI.

    To temporary enable autocomplete run `eval "$(pype autocomplete)"`.
    To install them, depending on your shell, run:
    - `bash`: `pype autocomplete --shell bash > ~/.bash_completion.d/pype`
    - `zsh` : `pype autocomplete --shell zsh  > ~/.zsh/completions/_pype`
    - `fish`: `pype autocomplete --shell fish > ~/.config/fish/completions/pype.fish`
    """
    pypeline_logger.log(f'Detected shell: "{shell}"')

    # Get the root command
    prog_name = get_current_root_name(ctx)
    pypeline_logger.debug_log(f'Program name: "{prog_name}"')
    complete_var = f"_{prog_name.upper()}_COMPLETE"
    pypeline_logger.debug_log(f'Complete variable: "{complete_var}"')

    # Requires Click 8.0+ for shell_complete
    # Create a completion context
    completion_cls = get_completion_class(shell)
    if completion_cls is None:
        click.echo(f'Shell "{shell}" is not supported', err=True)
        return

    completion = completion_cls(
        cli=ctx.find_root().command,
        ctx_args={},
        prog_name=prog_name,
        complete_var=complete_var,
    )
    click.echo(completion.source())


def main(args: Sequence[str]) -> None:
    # Divide click's argument from pipebox's arguments
    if "--" in args:
        position = args.index("--")
        click_args = args[:position]
        pipebox_args = args[position + 1 :]
    else:
        click_args = args
        pipebox_args = []

    # This is ok as click will pass the pipebox argument automatically
    pype(
        args=click_args,
        obj={
            "pipebox_args": pipebox_args,
        },
    )


def run():
    """Run the pipeline from the command line using the shell environment."""
    main(sys.argv[1:])


if __name__ == "__main__":
    run()