add

2026-10-15 20:36:59 +08:00
commit 6bf949949c
7923 changed files with 1369582 additions and 0 deletions
@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import sys
+from typing import Any, MutableMapping, cast
+
+import click
+from prompt_toolkit.history import InMemoryHistory
+
+from ._completer import ClickCompleter
+from .core import ReplContext
+from .exceptions import ClickExit  # type: ignore[attr-defined]
+from .exceptions import CommandLineParserError, ExitReplException, InvalidGroupFormat
+from .globals_ import ISATTY, get_current_repl_ctx
+from .utils import (
+    _execute_internal_and_sys_cmds,
+    _get_protected_args,
+    _set_protected_args,
+)
+
+__all__ = ["bootstrap_prompt", "register_repl", "repl"]
+
+
+def bootstrap_prompt(
+    group: click.MultiCommand,
+    prompt_kwargs: dict[str, Any],
+    ctx: click.Context,
+) -> dict[str, Any]:
+    """
+    Bootstrap prompt_toolkit kwargs or use user defined values.
+
+    :param group: click.MultiCommand object
+    :param prompt_kwargs: The user specified prompt kwargs.
+    """
+
+    defaults = {
+        "history": InMemoryHistory(),
+        "completer": ClickCompleter(group, ctx=ctx),
+        "message": "> ",
+    }
+
+    defaults.update(prompt_kwargs)
+    return defaults
+
+
+def repl(
+    old_ctx: click.Context,
+    prompt_kwargs: dict[str, Any] = {},
+    allow_system_commands: bool = True,
+    allow_internal_commands: bool = True,
+) -> None:
+    """
+    Start an interactive shell. All subcommands are available in it.
+
+    :param old_ctx: The current Click context.
+    :param prompt_kwargs: Parameters passed to
+        :py:func:`prompt_toolkit.PromptSession`.
+
+    If stdin is not a TTY, no prompt will be printed, but only commands read
+    from stdin.
+    """
+
+    group_ctx = old_ctx
+    # Switching to the parent context that has a Group as its command
+    # as a Group acts as a CLI for all of its subcommands
+    if old_ctx.parent is not None and not isinstance(
+        old_ctx.command, click.MultiCommand
+    ):
+        group_ctx = old_ctx.parent
+
+    group = cast(click.MultiCommand, group_ctx.command)
+
+    # An Optional click.Argument in the CLI Group, that has no value
+    # will consume the first word from the REPL input, causing issues in
+    # executing the command
+    # So, if there's an empty Optional Argument
+    for param in group.params:
+        if (
+            isinstance(param, click.Argument)
+            and group_ctx.params[param.name] is None  # type: ignore[index]
+            and not param.required
+        ):
+            raise InvalidGroupFormat(
+                f"{type(group).__name__} '{group.name}' requires value for "
+                f"an optional argument '{param.name}' in REPL mode"
+            )
+
+    # Delete the REPL command from those available, as we don't want to allow
+    # nesting REPLs (note: pass `None` to `pop` as we don't want to error if
+    # REPL command already not present for some reason).
+    repl_command_name = old_ctx.command.name
+
+    available_commands: MutableMapping[str, click.Command] = {}
+
+    if isinstance(group, click.CommandCollection):
+        available_commands = {
+            cmd_name: source.get_command(group_ctx, cmd_name)  # type: ignore[misc]
+            for source in group.sources
+            for cmd_name in source.list_commands(group_ctx)
+        }
+
+    elif isinstance(group, click.Group):
+        available_commands = group.commands
+
+    original_command = available_commands.pop(repl_command_name, None)  # type: ignore
+
+    repl_ctx = ReplContext(
+        group_ctx,
+        bootstrap_prompt(group, prompt_kwargs, group_ctx),
+        get_current_repl_ctx(silent=True),
+    )
+
+    if ISATTY:
+        # If stdin is a TTY, prompt the user for input using PromptSession.
+        def get_command() -> str:
+            return repl_ctx.session.prompt()  # type: ignore
+
+    else:
+        # If stdin is not a TTY, read input from stdin directly.
+        def get_command() -> str:
+            inp = sys.stdin.readline().strip()
+            repl_ctx._history.append(inp)
+            return inp
+
+    with repl_ctx:
+        while True:
+            try:
+                command = get_command()
+            except KeyboardInterrupt:
+                continue
+            except EOFError:
+                break
+
+            if not command:
+                if ISATTY:
+                    continue
+                else:
+                    break
+
+            try:
+                args = _execute_internal_and_sys_cmds(
+                    command, allow_internal_commands, allow_system_commands
+                )
+                if args is None:
+                    continue
+
+            except CommandLineParserError:
+                continue
+
+            except ExitReplException:
+                break
+
+            try:
+                # The group command will dispatch based on args.
+                old_protected_args = _get_protected_args(group_ctx)
+                try:
+                    _set_protected_args(group_ctx, args)
+                    group.invoke(group_ctx)
+                finally:
+                    _set_protected_args(group_ctx, old_protected_args)
+            except click.ClickException as e:
+                e.show()
+            except (ClickExit, SystemExit):
+                pass
+
+            except ExitReplException:
+                break
+
+    if original_command is not None:
+        available_commands[repl_command_name] = original_command  # type: ignore[index]
+
+
+def register_repl(group: click.Group, name="repl") -> None:
+    """Register :func:`repl()` as sub-command *name* of *group*."""
+    group.command(name=name)(click.pass_context(repl))