diff --git a/README.md b/README.md index 248e95f..e4802cd 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,12 @@ Install this plugin in the same environment as [LLM](https://llm.datasette.io/). ```bash llm install llm-cmd ``` + +To get syntax highlighting, install with +```bash +llm install 'llm-cmd[pygments]' +``` + ## Usage This command could be **very dangerous**. Do not use this unless you are confident you understand what it does and are sure you could spot if it is likely to do something dangerous. @@ -28,6 +34,12 @@ This will then be displayed in your terminal ready for you to edit it, or hit `< If the command doesnt't look right, hit `Ctrl+C` to cancel. +To enable syntax highlighting, use the `-H` flag with a Pygments style name, e.g. + +```bash +llm cmd -H monokai find all files in subdirectories that have a corresponding tilde backup +``` + ## The system prompt This is the prompt used by this tool: diff --git a/llm_cmd.py b/llm_cmd.py index 7d749cd..b283d10 100644 --- a/llm_cmd.py +++ b/llm_cmd.py @@ -4,7 +4,14 @@ from prompt_toolkit.shortcuts import prompt from prompt_toolkit.lexers import PygmentsLexer -from pygments.lexers.shell import BashLexer +from prompt_toolkit.styles.pygments import style_from_pygments_cls + +try: + from pygments.lexers.shell import BashLexer + from pygments.styles import get_style_by_name, get_all_styles + import pygments.util +except ImportError: + BashLexer = get_style_by_name = get_all_styles = None SYSTEM_PROMPT = """ @@ -24,11 +31,27 @@ def register_commands(cli): @click.argument("args", nargs=-1) @click.option("-m", "--model", default=None, help="Specify the model to use") @click.option("-s", "--system", help="Custom system prompt") + @click.option("-H", "--highlight-style", default=None, + help="Pygments highlight style, e.g. monokai") @click.option("--key", help="API key to use") - def cmd(args, model, system, key): + def cmd(args, model, system, key, highlight_style): """Generate and execute commands in your shell""" from llm.cli import get_default_model + style = None + if highlight_style is not None: + if get_style_by_name is None: + raise click.ClickException( + "Pygments is not installed, cannot use --highlight-style" + ) + try: + style = style_from_pygments_cls(get_style_by_name(highlight_style)) + except (ModuleNotFoundError, pygments.util.ClassNotFound): + raise click.ClickException( + f"Style {highlight_style} not found, available styles: " + f"{list(get_all_styles())}" + ) + prompt = " ".join(args) model_id = model or get_default_model() @@ -39,15 +62,23 @@ def cmd(args, model, system, key): result = model_obj.prompt(prompt, system=system or SYSTEM_PROMPT) - interactive_exec(str(result)) + interactive_exec(str(result), style) -def interactive_exec(command): +def interactive_exec(command, style): + if style is None: + kwargs = {} + else: + kwargs = { + "style": style, + "lexer": PygmentsLexer(BashLexer), + } if '\n' in command: print("Multiline command - Meta-Enter or Esc Enter to execute") - edited_command = prompt("> ", default=command, lexer=PygmentsLexer(BashLexer), multiline=True) - else: - edited_command = prompt("> ", default=command, lexer=PygmentsLexer(BashLexer)) + kwargs["multiline"] = True + + edited_command = prompt("> ", default=command, **kwargs) + try: output = subprocess.check_output( edited_command, shell=True, stderr=subprocess.STDOUT diff --git a/pyproject.toml b/pyproject.toml index cdb2ace..3969aaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ classifiers = [ dependencies = [ "llm", "prompt_toolkit>=3.0.43", - "pygments>=2.17.2", ] [project.urls] @@ -25,3 +24,4 @@ cmd = "llm_cmd" [project.optional-dependencies] test = ["pytest"] +pygments = ["pygments>=2.17.2"]