Files
shell_gpt/sgpt/utils.py
T
Farkhod Sadykov c48926a4a4 Non-interactive mode for --shell and integration changes (#455)
* Added --interaction that works with --shell option.
* Changed shell integrations to use new --no-interaction option.
* Moved shell integrations into dedicated file integration.py.
* Changed --install-integration logic to install integrations without downloading sh scripts.
* Removed validation for PROMPT argument, empty string by default.
* Fixing an issue when sgpt is being called from non-interactive shell environments.
* Fixed and optimised Dockerfile.
* README.md improvements, and new feature examples.
* New funny demo video.
2024-01-28 05:03:44 +01:00

96 lines
2.9 KiB
Python

import os
import platform
import shlex
from tempfile import NamedTemporaryFile
from typing import Any, Callable
import typer
from click import BadParameter, UsageError
from sgpt.__version__ import __version__
from sgpt.integration import bash_integration, zsh_integration
def get_edited_prompt() -> str:
"""
Opens the user's default editor to let them
input a prompt, and returns the edited text.
:return: String prompt.
"""
with NamedTemporaryFile(suffix=".txt", delete=False) as file:
# Create file and store path.
file_path = file.name
editor = os.environ.get("EDITOR", "vim")
# This will write text to file using $EDITOR.
os.system(f"{editor} {file_path}")
# Read file when editor is closed.
with open(file_path, "r", encoding="utf-8") as file:
output = file.read()
os.remove(file_path)
if not output:
raise BadParameter("Couldn't get valid PROMPT from $EDITOR")
return output
def run_command(command: str) -> None:
"""
Runs a command in the user's shell.
It is aware of the current user's $SHELL.
:param command: A shell command to run.
"""
if platform.system() == "Windows":
is_powershell = len(os.getenv("PSModulePath", "").split(os.pathsep)) >= 3
full_command = (
f'powershell.exe -Command "{command}"'
if is_powershell
else f'cmd.exe /c "{command}"'
)
else:
shell = os.environ.get("SHELL", "/bin/sh")
full_command = f"{shell} -c {shlex.quote(command)}"
os.system(full_command)
def option_callback(func: Callable) -> Callable: # type: ignore
def wrapper(cls: Any, value: str) -> None:
if not value:
return
func(cls, value)
raise typer.Exit()
return wrapper
@option_callback
def install_shell_integration(*_args: Any) -> None:
"""
Installs shell integration. Currently only supports ZSH and Bash.
Allows user to get shell completions in terminal by using hotkey.
Replaces current "buffer" of the shell with the completion.
"""
# TODO: Add support for Windows.
# TODO: Implement updates.
shell = os.getenv("SHELL", "")
if shell == "/bin/zsh":
typer.echo("Installing ZSH integration...")
with open(os.path.expanduser("~/.zshrc"), "a", encoding="utf-8") as file:
file.write(zsh_integration)
elif shell == "/bin/bash":
typer.echo("Installing Bash integration...")
with open(os.path.expanduser("~/.bashrc"), "a", encoding="utf-8") as file:
file.write(bash_integration)
else:
raise UsageError("ShellGPT integrations only available for ZSH and Bash.")
typer.echo("Done! Restart your shell to apply changes.")
@option_callback
def get_sgpt_version(*_args: Any) -> None:
"""
Displays the current installed version of ShellGPT
"""
typer.echo(f"ShellGPT {__version__}")