failprint ¤

def accept_custom_format(string: str) -> str:
    """Store the value in `formats` if it starts with custom.

    Arguments:
        string: A format name.

    Returns:
        The format name, or `custom` if it started with `custom=`.
    """
    if string.startswith("custom="):
        formats["custom"] = Format(string[7:])
        return "custom"
    return string

def add_flags(parser: ArgParser, *, set_defaults: bool = True) -> ArgParser:
    """Add some boolean flags to the parser.

    We made this method separate and public
    for its use in [duty](https://github.com/pawamoy/duty).

    Arguments:
        parser: The parser to add flags to.
        set_defaults: Whether to set default values on arguments.

    Returns:
        The augmented parser.
    """
    # IMPORTANT: the arguments destinations should match
    # the parameters names of the failprint.runners.run function.
    # As long as names are consistent between the two,
    # it's very easy to pass CLI args to the function,
    # and it also allows to avoid duplicating the parser arguments
    # in dependent projects like duty (https://github.com/pawamoy/duty) :)
    parser.add_argument(
        "-c",
        "--capture",
        choices=list(Capture),
        type=Capture,
        help="Which output to capture. Colors are supported with 'both' only, unless the command has a 'force color' option.",
    )
    parser.add_argument(
        "-f",
        "--fmt",
        "--format",
        dest="fmt",
        choices=formats.keys(),
        type=accept_custom_format,
        default=None,
        help="Output format. Pass your own Jinja2 template as a string with '-f custom=TEMPLATE'. "
        "Available variables: command, title (command or title passed with -t), code (exit status), "
        "success (boolean), failure (boolean), number (command number passed with -n), "
        "output (command output), nofail (boolean), quiet (boolean), silent (boolean). "
        "Available filters: indent (textwrap.indent).",
    )
    parser.add_bool_argument(
        ["-y", "--pty"],
        ["-Y", "--no-pty"],
        dest="pty",
        default=True if set_defaults else None,
        truthy_help="Enable the use of a pseudo-terminal. PTY doesn't allow programs to use standard input.",
        falsy_help="Disable the use of a pseudo-terminal. PTY doesn't allow programs to use standard input.",
    )
    parser.add_bool_argument(
        ["-p", "--progress"],
        ["-P", "--no-progress"],
        dest="progress",
        default=True if set_defaults else None,
        truthy_help="Print progress while running a command.",
        falsy_help="Don't print progress while running a command.",
    )
    # TODO: specific to the format
    parser.add_bool_argument(
        ["-q", "--quiet"],
        ["-Q", "--no-quiet"],
        dest="quiet",
        default=False if set_defaults else None,
        truthy_help="Don't print the command output, even if it failed.",
        falsy_help="Print the command output when it fails.",
    )
    # TODO: specific to the format
    parser.add_bool_argument(
        ["-s", "--silent"],
        ["-S", "--no-silent"],
        dest="silent",
        default=False if set_defaults else None,
        truthy_help="Don't print anything.",
        falsy_help="Print output as usual.",
    )
    parser.add_bool_argument(
        ["-z", "--zero", "--nofail"],
        ["-Z", "--no-zero", "--strict"],
        dest="nofail",
        default=False if set_defaults else None,
        truthy_help="Don't fail. Always return a success (0) exit code.",
        falsy_help="Return the original exit code.",
    )
    return parser

def as_python_statement(func: Callable | LazyCallable, args: Sequence | None = None, kwargs: dict | None = None) -> str:
    """Transform a callable and its arguments into a Python statement string.

    Arguments:
        func: The callable to transform.
        args: Positional arguments passed to the function.
        kwargs: Keyword arguments passed to the function.

    Returns:
        A Python statement.
    """
    if isinstance(func, LazyCallable):
        callable_name = func.name or _get_callable_name(func.call)
        args = args or func.args
        kwargs = kwargs or func.kwargs
    else:
        callable_name = _get_callable_name(func)
    args_str = [repr(arg) for arg in args] if args else []
    kwargs_str = [f"{k}={v!r}" for k, v in kwargs.items()] if kwargs else []
    arguments = ", ".join(args_str + kwargs_str)
    return f"{callable_name}({arguments})"

def as_shell_command(cmd: list[str]) -> str:
    """Rebuild a command line from system arguments.

    Arguments:
        cmd: The command as a list of strings.

    Returns:
        A printable and shell-runnable command.
    """
    parts = []
    for part in cmd:
        if not part:
            parts.append('""')
            continue
        has_spaces = " " in part
        has_double_quotes = '"' in part
        has_single_quotes = "'" in part
        if has_double_quotes and not has_single_quotes:
            # double quotes, no single quotes
            # -> wrap in single quotes
            part = f"'{part}'"  # noqa: PLW2901
        elif has_single_quotes and has_double_quotes:
            # double and single quotes
            # -> escape double quotes, wrap in double quotes
            part = part.replace('"', r"\"")  # noqa: PLW2901
            part = f'"{part}"'  # noqa: PLW2901
        elif has_single_quotes or has_spaces:
            # spaces or single quotes
            # -> wrap in double quotes
            part = f'"{part}"'  # noqa: PLW2901
        parts.append(part)
    return " ".join(parts)

def get_parser() -> ArgParser:
    """Return the CLI argument parser.

    Returns:
        An argparse parser.
    """
    parser = add_flags(ArgParser(prog="failprint"))
    # TODO: specific to the format
    parser.add_argument("-n", "--number", type=int, default=1, help="Command number. Useful for the 'tap' format.")
    # TODO: specific to the format
    parser.add_argument("-t", "--title", help="Command title. Default is the command itself.")
    parser.add_argument("cmd", metavar="COMMAND", nargs="+")
    parser.add_argument("-V", "--version", action="version", version=f"%(prog)s {debug._get_version()}")
    parser.add_argument("--debug-info", action=_DebugInfo, help="Print debug information.")
    return parser

def main(args: list[str] | None = None) -> int:
    """Run the main program.

    This function is executed when you type `failprint` or `python -m failprint`.

    Parameters:
        args: Arguments passed from the command line.

    Returns:
        An exit code.
    """
    parser = get_parser()
    opts = parser.parse_args(args).__dict__.items()
    return run(**{_: value for _, value in opts if value is not None}).code

def printable_command(cmd: CmdFuncType, args: Sequence | None = None, kwargs: dict | None = None) -> str:
    """Transform a command or function into a string.

    Arguments:
        cmd: The command or function to transform.
        args: Positional arguments passed to the function.
        kwargs: Keyword arguments passed to the function.

    Returns:
        A shell command or python statement string.
    """
    if isinstance(cmd, str):
        return cmd
    if callable(cmd):
        return as_python_statement(cmd, args, kwargs)
    return as_shell_command(cmd)

def run(
    cmd: CmdFuncType,
    *,
    args: Sequence | None = None,
    kwargs: dict | None = None,
    number: int = 1,
    capture: str | bool | Capture | None = None,
    title: str | None = None,
    fmt: str | None = None,
    pty: bool = False,
    progress: bool = True,
    nofail: bool = False,
    quiet: bool = False,
    silent: bool = False,
    stdin: str | None = None,
    command: str | None = None,
) -> RunResult:
    """Run a command in a subprocess or a Python function, and print its output if it fails.

    Arguments:
        cmd: The command to run.
        args: Arguments to pass to the callable.
        kwargs: Keyword arguments to pass to the callable.
        number: The command number.
        capture: The output to capture.
        title: The command title.
        fmt: The output format.
        pty: Whether to run in a PTY.
        progress: Whether to show progress.
        nofail: Whether to always succeed.
        quiet: Whether to not print the command output.
        silent: Don't print anything.
        stdin: String to use as standard input.
        command: The command to display.

    Returns:
        The command exit code, or 0 if `nofail` is True.
    """
    format_name: str = fmt or os.environ.get("FAILPRINT_FORMAT", DEFAULT_FORMAT)  # type: ignore[assignment]
    format_name = accept_custom_format(format_name)
    format_obj = formats.get(format_name, formats[DEFAULT_FORMAT])

    env = Environment(autoescape=False)  # noqa: S701 (no HTML: no need to escape)
    env.filters["indent"] = textwrap.indent
    env.filters["escape"] = env.filters["e"] = escape
    env.filters["unescape"] = env.filters["u"] = unescape

    command = command if command is not None else printable_command(cmd, args, kwargs)

    if not silent and progress and format_obj.progress_template:
        progress_template = env.from_string(format_obj.progress_template)
        print(unescape(parse(progress_template.render({"title": title, "command": command}))), end="\r")  # noqa: T201

    capture = Capture.cast(capture)

    if callable(cmd):
        code, output = run_function(cmd, args=args, kwargs=kwargs, capture=capture, stdin=stdin)
    else:
        code, output = run_command(cmd, capture=capture, ansi=format_obj.accept_ansi, pty=pty, stdin=stdin)

    if not silent:
        template = env.from_string(format_obj.template)
        rendered = template.render(
            {
                "title": title,
                "command": command,
                "code": code,
                "success": code == 0,
                "failure": code != 0,
                "number": number,
                "output": output,
                "nofail": nofail,
                "quiet": quiet,
                "silent": silent,
            },
        )
        print(unescape(parse(rendered)))  # noqa: T201

    return RunResult(0 if nofail else code, output)

def run_command(
    cmd: CmdType,
    *,
    capture: Capture = Capture.BOTH,
    ansi: bool = False,
    pty: bool = False,
    stdin: str | None = None,
) -> tuple[int, str]:
    """Run a command.

    Arguments:
        cmd: The command to run.
        capture: The output to capture.
        ansi: Whether to accept ANSI sequences.
        pty: Whether to run in a PTY.
        stdin: String to use as standard input.

    Returns:
        The exit code and the command output.
    """
    shell = isinstance(cmd, str)

    # if chosen format doesn't accept ansi, or on Windows, don't use pty
    if pty and (not ansi or WINDOWS):
        pty = False

    # pty can only combine, so only use pty when combining
    if pty and capture in {Capture.BOTH, Capture.NONE}:
        if shell:
            cmd = ["sh", "-c", cmd]  # type: ignore[list-item]  # we know cmd is str
        return run_pty_subprocess(cmd, capture=capture, stdin=stdin)  # type: ignore[arg-type]  # we made sure cmd is a list

    # we are on Windows
    if WINDOWS:
        # make sure the process can find the executable
        if not shell:
            cmd[0] = shutil.which(cmd[0]) or cmd[0]  # type: ignore[index]  # we know cmd is a list
        return run_subprocess(cmd, capture=capture, shell=shell, stdin=stdin)

    return run_subprocess(cmd, capture=capture, shell=shell, stdin=stdin)

def run_function(
    func: Callable,
    *,
    args: Sequence | None = None,
    kwargs: dict | None = None,
    capture: Capture = Capture.BOTH,
    stdin: str | None = None,
) -> tuple[int, str]:
    """Run a function.

    Arguments:
        func: The function to run.
        args: Positional arguments passed to the function.
        kwargs: Keyword arguments passed to the function.
        capture: The output to capture.
        stdin: String to use as standard input.

    Returns:
        The exit code and the function output.
    """
    args = args or []
    kwargs = kwargs or {}

    if capture == Capture.NONE:
        return run_function_get_code(func, args=args, kwargs=kwargs), ""

    with capture.here(stdin=stdin) as captured:
        code = run_function_get_code(func, args=args, kwargs=kwargs)

    return code, str(captured)

def run_function_get_code(
    func: Callable,
    *,
    args: Sequence,
    kwargs: dict,
) -> int:
    """Run a function and return a exit code.

    Arguments:
        func: The function to run.
        args: Positional arguments passed to the function.
        kwargs: Keyword arguments passed to the function.

    Returns:
        An exit code.
    """
    try:
        result = func(*args, **kwargs)
    except SystemExit as exit:
        if exit.code is None:
            return 0
        if isinstance(exit.code, int):
            return exit.code
        sys.stderr.write(str(exit.code))
        return 1
    except Exception:  # noqa: BLE001
        sys.stderr.write(traceback.format_exc() + "\n")
        return 1

    # if func was a lazy callable, recurse
    if isinstance(result, LazyCallable):
        return run_function_get_code(result, args=(), kwargs={})

    # first check True and False
    # because int(True) == 1 and int(False) == 0
    if result is True:
        return 0
    if result is False:
        return 1
    try:
        return int(result)
    except (ValueError, TypeError):
        if result is None or bool(result):
            return 0
        return 1

def run_pty_subprocess(
    cmd: list[str],
    *,
    capture: Capture = Capture.BOTH,
    stdin: str | None = None,
) -> tuple[int, str]:
    """Run a command in a PTY subprocess.

    Arguments:
        cmd: The command to run.
        capture: The output to capture.
        stdin: String to use as standard input.

    Returns:
        The exit code and the command output.
    """
    process = PtyProcessUnicode.spawn(cmd)
    process.delayafterclose = 0.01  # default to 0.1
    process.delayafterterminate = 0.01  # default to 0.1
    pty_output: list[str] = []

    if stdin is not None:
        process.setecho(state=False)
        process.waitnoecho()
        process.write(stdin)
        process.sendeof()
        # not sure why but sending only one eof is not always enough,
        # so we send a second one and ignore any IO error
        with contextlib.suppress(OSError):
            process.sendeof()

    while True:
        try:
            output_data = process.read()
        except EOFError:
            break
        if capture == Capture.NONE:
            print(output_data, end="", flush=True)  # noqa: T201
        else:
            pty_output.append(output_data)

    output = "".join(pty_output).replace("\r\n", "\n")
    return process.wait(), output

def run_subprocess(
    cmd: CmdType,
    *,
    capture: Capture = Capture.BOTH,
    shell: bool = False,
    stdin: str | None = None,
) -> tuple[int, str]:
    """Run a command in a subprocess.

    Arguments:
        cmd: The command to run.
        capture: The output to capture.
        shell: Whether to run the command in a shell.
        stdin: String to use as standard input.

    Returns:
        The exit code and the command raw output.
    """
    if capture == Capture.NONE:
        stdout_opt = None
        stderr_opt = None
    else:
        stdout_opt = subprocess.PIPE
        stderr_opt = subprocess.STDOUT if capture == Capture.BOTH else subprocess.PIPE

    if shell and not isinstance(cmd, str):
        cmd = printable_command(cmd)

    process = subprocess.run(  # noqa: S603
        cmd,
        input=stdin,
        stdout=stdout_opt,
        stderr=stderr_opt,
        shell=shell,
        text=True,
        encoding="utf8",
        check=False,
    )

    if capture == Capture.NONE:
        output = ""
    elif capture == Capture.STDERR:
        output = process.stderr
    else:
        output = process.stdout

    return process.returncode, output