Make pygments dependency optional

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:

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

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"]