pyxargs

#!/usr/bin/python3.11

import re
import subprocess as shell
import sys
import os
import json
import tempfile

from argparse import ArgumentParser

class Shell:
    def __init__(self, cmds, timeout=1, tries=5):
        self.cmds = cmds
        self.procs = {}
        self.status = {}
        self.timeout = timeout
        self.tries = tries

    def run(self) -> None:
        for cmd in self.cmds:
            pipe = shell.PIPE
            self.procs[cmd] = shell.Popen([cmd], shell=True, stdout=pipe, stderr=pipe)

    def is_complete(self):
        returned = [x for x in self.status.keys() if self.status[x]["returncode"]]
        return len(returned) == len(self.cmds)

    def update_status(self, cmd, prop, value):
        if not self.status.get(cmd):
            self.status[cmd] = {"returncode": False, "error": False, "output": False}

        self.status[cmd][prop] = value

    def wait_once(self, cmds: list) -> bool:
        while len(cmds) > 0:
            current = cmds[0]
            cmds = cmds[1:]
            proc = self.procs[current]
            timeout = self.timeout

            try:
                out, errs = proc.communicate(timeout=timeout)
                self.update_status(current, "returncode", proc.returncode)

                if proc.returncode == 0:
                    self.update_status(
                        current, "output", out.decode().strip().split("\n")
                    )
                elif errs:
                    errs = errs.decode().strip()
                    if len(errs) > 0:
                        self.update_status(
                            current, "error", errs.split("\n")
                        )
                    else:
                        self.update_status(current, "error", False)

                    self.update_status(current, "output", False)

            except shell.TimeoutExpired:
                pass

        return self.is_complete()

    def wait(self) -> bool:
        tries = self.tries
        cmds = self.cmds
        current_try = 0

        while current_try != tries:
            if self.wait_once(cmds):
                break
            else:
                cmds = [x for x in cmds if not self.status.get(x)]
                current_try = current_try + 1

        return self.is_complete()

    @staticmethod
    def main(cmds, timeout=1, tries=5):
        runner = Shell(cmds, timeout=timeout, tries=tries)
        runner.run()
        runner.wait()

        return runner.status


def default_placeholders(x: str) -> dict:
    basename = os.path.basename(x)
    abspath = os.path.abspath(x)
    dirname = os.path.dirname(x)

    return dict(basename=basename, abspath=abspath, dirname=dirname)


def add_default_subs(xs: list) -> list:
    out = []

    for x in xs:
        out.append(x.format(**default_placeholders(x)))

    return out


def write_output(output: bytes):
    with tempfile.NamedTemporaryFile(delete=False) as fh:
        fh.write(output)
        print(f"output written to", fh.name)


def parse_cmdline() -> tuple[tuple, dict]:
    parser = ArgumentParser("python equivalent for xargs")
    add = parser.add_argument

    add(
        "-x",
        "--cmd",
        nargs="+",
        required=True,
        help="command with python placeholder strings",
    )

    add(
        "-t",
        "--timeout",
        default=10,
        help="default timeout to check for polling for output",
    )

    add("-p", "--path", help="read lines for path")
    add("-i", "--stdin", action="store_true", help="read lines from stdin (default)")
    add("-/" ,"--replace", nargs="+", help="placeholder definitions")
    add("-d", "--sep", nargs=1, help='record separation (default: "\n")', default="\n")
    add("--stdout", action="store_true", help="print output as json")

    add(
        "-o",
        "--dest",
        nargs=1,
        help="save output to path",
        metavar="PATH",
        default=False,
    )

    args: dict = vars(parser.parse_args())
    definitions: list = args.get("defs", [])
    defs: dict = {}
    l = len(definitions)
    path: str = args["path"]
    sep: str = args["sep"][0]
    records: list[str] = []
    timeout: int = int(args["timeout"])
    cmd: str = (" ").join(args["cmd"])
    cmds: list[str] = []
    failed = []

    if path:
        if not os.path.isfile(path):
            raise ValueError(f"invalid path {path}")
        else:
            with open(path) as fh:
                records = fh.read().split(sep)
    elif not args["stdin"]:
        raise ValueError(f"no records provided. use -i or pass records through stdin")
    else:
        records = sys.stdin.read().split(sep)

    if l == 0:
        cmds = add_default_subs(records)
    elif l % 2 != 0:
        raise ValueError(f"expected an even number of args, got {l}")
    else:

        def parse_cmds(cmd: str, recs: list, defs: dict) -> tuple[list, list]:
            out = []

            for i, x in enumerate(recs):
                found = default_placeholders(x)
                items = defs.items()
                fails = 0

                for var, r in items:
                    r = re.compile(r)

                    if m := re.search(r, x):
                        found[var] = m.group()
                    else:
                        fails = fails + 1

                if fails != len(items):
                    out.append(cmd.format(**found))
                else:
                    failed.append((i, x))

            return (out, failed)

        placeholder: list = [definitions[i] for i in range(l) if i % 2 == 0]
        sub: list = [definitions[i] for i in range(l) if i % 2 != 0]
        defs = dict(zip(placeholder, sub))
        (cmds, failed) = parse_cmds(cmd, records, defs)

    return (
        (cmds, failed),
        {"timeout": timeout, "stdout": args["stdout"], "dest": args["dest"]},
    )


def main():
    [cmds, failed], opts = parse_cmdline()
    status = Shell.main(cmds, timeout=opts["timeout"])
    status = json.dumps(status)

    if opts["stdout"]:
        print(status)
    elif opts["dest"]:
        with open(opts["dest"]) as fh:
            fh.write(status)
            print(f"output written to", opts["dest"])
    else:
        write_output(status.encode())


main()