feat: copy to clipboard

gptme/clipboard.py

@@ -0,0 +1,45 @@
+import platform
+import subprocess
+
+
+text = ""
+
+
+def set_copytext(new_text: str):
+    global text
+    text = new_text
+
+
+def copy() -> bool:
+    """return True if successful"""
+    from .util import get_installed_programs
+
+    global text
+    if platform.system() == "Linux":
+        # check if xclip or wl-clipboard is installed
+        installed = get_installed_programs(("xclip", "wl-copy"))
+        if "wl-copy" in installed:
+            output = subprocess.run(["wl-copy"], input=text, text=True, check=True)
+            if output.returncode != 0:
+                print("wl-copy failed to copy to clipboard.")
+                return False
+            return True
+        elif "xclip" in installed:
+            output = subprocess.run(
+                ["xclip", "-selection", "clipboard"], input=text, text=True
+            )
+            if output.returncode != 0:
+                print("xclip failed to copy to clipboard.")
+                return False
+            return True
+        else:
+            print("No clipboard utility found. Please install xclip or wl-clipboard.")
+            return False
+    elif platform.system() == "Darwin":
+        output = subprocess.run(["pbcopy"], input=text, text=True)
+        if output.returncode != 0:
+            print("pbcopy failed to copy to clipboard.")
+            return False
+        return True
+
+    return False

gptme/tools/shell.py

@@ -3,47 +3,40 @@
 """
 
 import atexit
-import functools
 import logging
 import os
 import re
 import select
-import shutil
 import subprocess
 import sys
 from collections.abc import Generator
 
 import bashlex
 
 from ..message import Message
-from ..util import get_tokenizer, print_preview
+from ..util import get_tokenizer, print_preview, get_installed_programs
 from .base import ConfirmFunc, ToolSpec, ToolUse
 
 logger = logging.getLogger(__name__)
 
 
-@functools.lru_cache
-def get_installed_programs() -> set[str]:
-    candidates = [
-        # platform-specific
-        "brew",
-        "apt-get",
-        "pacman",
-        # common and useful
-        "ffmpeg",
-        "magick",
-        "pandoc",
-        "git",
-        "docker",
-    ]
-    installed = set()
-    for candidate in candidates:
-        if shutil.which(candidate) is not None:
-            installed.add(candidate)
-    return installed
-
-
-shell_programs_str = "\n".join(f"- {prog}" for prog in get_installed_programs())
+candidates = (
+    # platform-specific
+    "brew",
+    "apt-get",
+    "pacman",
+    # common and useful
+    "ffmpeg",
+    "magick",
+    "pandoc",
+    "git",
+    "docker",
+)
+
+
+shell_programs_str = "\n".join(
+    f"- {prog}" for prog in get_installed_programs(candidates)
+)
 is_macos = sys.platform == "darwin"
 
 instructions = f"""
@@ -260,7 +253,7 @@ def execute_shell(
                 break
 
     if not allowlisted:
-        print_preview(cmd, "bash")
+        print_preview(cmd, "bash", True)
         if not confirm("Run command?"):
             yield Message("system", "User chose not to run command.")
             return

gptme/tools/tmux.py

@@ -157,7 +157,7 @@ def execute_tmux(
     assert not args
     cmd = code.strip()
 
-    print_preview(f"Command: {cmd}", "bash")
+    print_preview(f"Command: {cmd}", "bash", copy=True)
     if not confirm(f"Execute command: {cmd}?"):
         yield Message("system", "Command execution cancelled.")
         return

gptme/util.py

@@ -1,6 +1,8 @@
 import io
 import logging
 import random
+import shutil
+import functools
 import re
 import sys
 import termios
@@ -10,6 +12,8 @@
 from pathlib import Path
 from typing import Any
 
+from .clipboard import copy, set_copytext
+
 import tiktoken
 from rich import print
 from rich.console import Console
@@ -132,9 +136,27 @@ def epoch_to_age(epoch, incl_date=False):
         )
 
 
-def print_preview(code: str, lang: str):  # pragma: no cover
+copiable = False
+
+
+def set_copiable():
+    global copiable
+    copiable = True
+
+
+def clear_copiable():
+    global copiable
+    copiable = False
+
+
+def print_preview(code: str, lang: str, copy: bool = False):  # pragma: no cover
     print()
     print("[bold white]Preview[/bold white]")
+
+    if copy:
+        set_copiable()
+        set_copytext(code)
+
     # NOTE: we can set background_color="default" to remove background
     print(Syntax(code.strip("\n"), lang))
     print()
@@ -148,10 +170,19 @@ def ask_execute(question="Execute code?", default=True) -> bool:  # pragma: no c
     termios.tcflush(sys.stdin, termios.TCIFLUSH)  # flush stdin
 
     choicestr = f"[{'Y' if default else 'y'}/{'n' if default else 'N'}]"
+    copystr = r"\[c] to copy" if copiable else ""
     answer = console.input(
-        f"[bold bright_yellow on red] {question} {choicestr} [/] ",
+        f"[bold bright_yellow on red] {question} {choicestr}{copystr} [/] ",
     )
+
     global override_auto
+
+    if not override_auto and copiable and "c" == answer.lower().strip():
+        if copy():
+            print("Copied to clipboard.")
+            return False
+        clear_copiable()
+
     if answer.lower() in [
         "auto"
     ]:  # secret option to stop asking for the rest of the session
@@ -309,3 +340,12 @@ def path_with_tilde(path: Path) -> str:
     if path_str.startswith(home):
         return path_str.replace(home, "~", 1)
     return path_str
+
+
+@functools.lru_cache
+def get_installed_programs(candidates: tuple[str, ...]) -> set[str]:
+    installed = set()
+    for candidate in candidates:
+        if shutil.which(candidate) is not None:
+            installed.add(candidate)
+    return installed