run_shell_cmd.py

"""
This module provides utilities to run shell commands with enhanced features such as:

- Automatic detection of the operating system prompt symbol.
- Colored command prompt and output using colorama (with fallback if colorama is
  not installed).
- Flattening of command arguments that can be strings or sequences of strings.
- Automatic prepending of 'uv run' to Python commands if the 'uv' tool is installed.
- Optional working directory and environment variable support.
- Capturing and displaying command output with customizable colors.

Functions:
- get_prompt_symbol(): Returns the shell prompt symbol based on the OS.
- color_text(text, color): Wraps text with colorama color codes if color is provided.
- is_uv_installed(): Checks if the 'uv' command is available in the system PATH.
- run_command(*args, cwd, env, prompt_color, stdout_color, stderr_color,
  show_output): Runs a shell command with colored output and optional environment
  and working directory.

Example usage:
    run_command("python", "--version")
    run_command("python", ["-c", "print('Hello from list arg')"])
    run_command("python", ["-c", "print('Hello')"], "--version")
"""

import os
import subprocess
import sys
import shutil
from typing import Optional, Dict, Union, Sequence, Iterator

try:
    from colorama import Fore, Style, init
    init(autoreset=True)
except ImportError:
    # Define dummy Fore and Style with empty strings
    class Dummy:
        RESET_ALL = ''
        RED = ''
        GREEN = ''
        YELLOW = ''
        BLUE = ''
        MAGENTA = ''
        CYAN = ''
        WHITE = ''
        # Add other colors if you use them

    Fore = Dummy()
    Style = Dummy()

PROMPT_SYMBOLS = {
    "nt": ">",
    "posix": "$"
}

def get_prompt_symbol() -> str:
    """Return the prompt symbol based on the current OS."""
    return PROMPT_SYMBOLS.get(os.name, "$")

def color_text(text: str, color: Optional[str] = None) -> str:
    """Wrap text with colorama color codes if color is provided."""
    return f"{color}{text}{Style.RESET_ALL}" if color else text

def is_uv_installed() -> bool:
    """Check if 'uv' command is available in the system PATH."""
    return shutil.which("uv") is not None

def _flatten_args(args: Sequence[Union[str, Sequence[str]]]) -> Iterator[str]:
    """
    Flatten a sequence of arguments where each argument can be a string or a
    sequence of strings.

    Yields individual string arguments.
    """
    for arg in args:
        if isinstance(arg, str):
            yield arg
        elif isinstance(arg, (list, tuple)):
            for sub_arg in arg:
                if not isinstance(sub_arg, str):
                    raise TypeError(f"Nested argument {sub_arg} is not a string")
                yield sub_arg
        else:
            raise TypeError(f"Argument {arg} is not a string or sequence of strings")

def run_command(
    *args: Union[str, Sequence[str]],
    cwd: Optional[str] = None,
    env: Optional[Dict[str, str]] = None,
    prompt_color: str = Fore.CYAN,
    stdout_color: str = Fore.GREEN,
    stderr_color: str = Fore.RED,
    show_output: bool = True
) -> Optional[subprocess.CompletedProcess]:
    """
    Run a shell command with colored output and optional environment and working
    directory.

    Accepts multiple command arguments directly. Each argument can be a string or
    a sequence of strings, which will be flattened into a single list of command
    arguments.

    If the first argument is 'python' and 'uv' is installed, it prepends 'uv run'
    to the command.

    Args:
        *args: Command and its arguments as strings or sequences of strings.
        cwd: Optional working directory to run the command in.
        env: Optional environment variables dictionary.
        prompt_color: Color for the command prompt.
        stdout_color: Color for standard output.
        stderr_color: Color for standard error.
        show_output: Whether to print the command output.

    Returns:
        subprocess.CompletedProcess instance or None if an error occurred.
    """
    if not args:
        print(color_text("No command specified.", stderr_color), file=sys.stderr)
        return None

    # Flatten the args into a single list of strings
    try:
        command_args = list(_flatten_args(args))
    except TypeError as e:
        print(color_text(f"Invalid argument: {e}", stderr_color), file=sys.stderr)
        return None

    if not command_args:
        print(color_text("No command arguments after flattening.", stderr_color),
              file=sys.stderr)
        return None

    # Ensure environment variables are strings
    if env:
        env = {str(k): str(v) for k, v in env.items()}

    # Prepend 'uv run' if applicable
    if command_args[0] == "python" and is_uv_installed():
        command_args = ["uv", "run"] + command_args

    prompt = get_prompt_symbol()
    import shlex
    command_display = " ".join(shlex.quote(arg) for arg in command_args)
    print(f"\n{prompt} {color_text(command_display, prompt_color)}")

    try:
        result = subprocess.run(
            command_args,
            shell=False,
            cwd=cwd,
            env=env,
            text=True,
            encoding="utf-8",
            capture_output=True,
            check=False
        )
    except Exception as exc:
        print(color_text(f"Error running command: {exc}", stderr_color),
              file=sys.stderr)
        return None

    if show_output:
        if result.stdout and result.stdout.strip():
            print(color_text(result.stdout, stdout_color), end="")
        if result.stderr and result.stderr.strip():
            print(color_text(result.stderr, stderr_color), file=sys.stderr, end="")

    return result


if __name__ == "__main__":
    # Examples:

    # Simple command with multiple string args
    run_command("python", "--version")

    # Command with a list argument that will be flattened
    run_command("python", ["-c", "print('Hello from list arg')"])

    # Mixed string and sequence arguments
    run_command("python", ["-c", "print('Hello')"], "--version")