iripau 0.1.0__py3-none-any.whl

iripau/subprocess.py

@@ -0,0 +1,526 @@
+"""
+A wrapper of the subprocess module
+
+This module relies on the following system utilities being installed:
+* bash
+* kill
+* pstree
+* sudo
+* tee
+"""
+
+import io
+import os
+import re
+import sys
+import shlex
+import psutil
+import subprocess
+
+from subprocess import DEVNULL
+from subprocess import PIPE
+from subprocess import STDOUT
+from subprocess import CompletedProcess
+from subprocess import TimeoutExpired
+from subprocess import SubprocessError  # noqa: F401
+from subprocess import CalledProcessError
+
+from time import time
+from typing import Union, Iterable, Callable
+from tempfile import SpooledTemporaryFile
+from contextlib import contextmanager, nullcontext
+
+FILE = -4
+GLOBAL_ECHO = False
+GLOBAL_STDOUTS = set()
+GLOBAL_STDERRS = set()
+GLOBAL_PROMPTS = set()
+
+
+TeeStream = Union[io.IOBase, Callable[[], io.IOBase]]
+
+
+class PipeFile(SpooledTemporaryFile):
+    """ A file to be used as stdin, stdout and stderr in Popen to avoid dead lock
+        when the process output is too long using PIPE.
+
+        If used as stdin, the content should be written before spawning the process.
+    """
+
+    def __init__(self, content=None, encoding=None, errors=None, text=None):
+        super().__init__(
+            mode="w+t" if text else "w+b",
+            encoding=encoding,
+            errors=errors
+        )
+
+        if content:
+            self.write(content)
+            self.seek(0)
+
+    def read_all(self):
+        self.seek(0)
+        return self.read()
+
+
+class Tee(subprocess.Popen):
+    """ A subprocess to send real-time input to several file descriptors """
+
+    def __init__(self, input, fds, output=None, encoding=None, errors=None, text=None):
+        if output is STDOUT:
+            raise ValueError("output cannot be STDOUT")
+
+        if output is None:
+            output = DEVNULL
+
+        fds = normalize_outerr_fds(fds)
+        if 1 in fds:
+            if output != DEVNULL:
+                fds.add(2)
+            stdout = None
+            stderr = output
+        elif 2 in fds:
+            stdout = output
+            stderr = None
+        else:
+            stdout = output
+            stderr = DEVNULL
+
+        fds.discard(1)
+        super().__init__(
+            stdin=input, stdout=stdout, stderr=stderr, pass_fds=fds - {2},
+            encoding=encoding, errors=errors, text=text, **self.get_kwargs(fds)
+        )
+
+        self.output = self.stdout if self.stderr is None else self.stderr
+
+    @staticmethod
+    def get_cmd(fds):
+        return ["tee", "-a"] + [f"/dev/fd/{fd}" for fd in fds]
+
+    if os.access("/dev/fd/2", os.W_OK):
+        @classmethod
+        def get_kwargs(cls, fds):
+            return {"args": cls.get_cmd(fds)}
+    else:
+        @classmethod
+        def get_kwargs(cls, fds):
+            if 2 in fds:
+                return {
+                    "args": " ".join(cls.get_cmd(fds - {2})) + " >(cat >&2)",
+                    "shell": True
+                }
+            return {"args": cls.get_cmd(fds)}
+
+    def communicate(self, *args, **kwargs):
+        stdout, stderr = super().communicate(*args, **kwargs)
+        return stdout if stderr is None else stderr
+
+
+class Popen(subprocess.Popen):
+    """ A subprocess.Popen that can send its stdout and stderr to several files
+        in real-time keeping the ability of capturing its output.
+    """
+
+    def __init__(
+        self, args, *, cwd=None, env=None, encoding=None, errors=None, text=None,
+        stdout_tees: Iterable[TeeStream] = [], add_global_stdout_tees=True,
+        stderr_tees: Iterable[TeeStream] = [], add_global_stderr_tees=True,
+        prompt_tees: Iterable[TeeStream] = [], add_global_prompt_tees=True,
+        echo=None, alias=None, comment=None, **kwargs
+    ):
+        stdout = kwargs.get("stdout")
+        stderr = kwargs.get("stderr")
+
+        stdout_tees, stderr_tees, prompt_tees = self._get_tee_sets(
+            stdout_tees, add_global_stdout_tees,
+            stderr_tees, add_global_stderr_tees,
+            prompt_tees, add_global_prompt_tees,
+            echo, stdout, stderr
+        )
+
+        if stderr is STDOUT:
+            err2out = True
+            stderr_tees = set()
+        else:
+            err2out = False
+
+        stdout_fds = {tee.fileno() for tee in stdout_tees}
+        stderr_fds = {tee.fileno() for tee in stderr_tees}
+        prompt_fds = {tee.fileno() for tee in prompt_tees}
+
+        if stdout_fds:
+            kwargs["stdout"] = PIPE
+
+        if stderr_fds:
+            kwargs["stderr"] = PIPE
+
+        if prompt_fds:
+            stream_prompts(prompt_fds, alias or args, cwd, env, err2out, comment)
+
+        super().__init__(args, cwd=cwd, env=env,
+                         encoding=encoding, errors=errors, text=text, **kwargs)
+
+        self.original_stdout = self.stdout
+        self.original_stderr = self.stderr
+
+        stdout_process = stderr_process = self
+        if stdout_fds:
+            stdout_process = Tee(self.stdout, stdout_fds, stdout, encoding, errors, text)
+            self.stdout = stdout_process.output
+        if stderr_fds:
+            stderr_process = Tee(self.stderr, stderr_fds, stderr, encoding, errors, text)
+            self.stderr = stderr_process.output
+
+        self.stdout_process = stdout_process
+        self.stderr_process = stderr_process
+
+    @staticmethod
+    def _get_tee_files(tees: Iterable[TeeStream]):
+        return set(callable(tee) and tee() or tee for tee in tees)
+
+    @classmethod
+    def _get_tee_sets(
+        cls,
+        stdout_tees, add_global_stdout_tees,
+        stderr_tees, add_global_stderr_tees,
+        prompt_tees, add_global_prompt_tees,
+        echo, stdout, stderr
+    ):
+        stdout_tees = set(stdout_tees)
+        stderr_tees = set(stderr_tees)
+        prompt_tees = set(prompt_tees)
+
+        if add_global_prompt_tees:
+            prompt_tees.update(GLOBAL_PROMPTS)
+        if add_global_stdout_tees:
+            stdout_tees.update(GLOBAL_STDOUTS)
+        if add_global_stderr_tees:
+            stderr_tees.update(GLOBAL_STDERRS)
+
+        if echo is None:
+            echo = GLOBAL_ECHO
+
+        if echo:
+            prompt_tees.add(sys.stdout)
+            stdout_tees.add(sys.stdout)
+            stderr_tees.add(sys.stderr)
+
+        if stdout is None:
+            if stdout_tees == {sys.stdout}:
+                stdout_tees.clear()  # tee process not needed for stdout
+            if stdout_tees:
+                stdout_tees.add(sys.stdout)
+
+        if stderr is None:
+            if stderr_tees == {sys.stderr}:
+                stderr_tees.clear()  # tee process not needed for stderr
+            if stderr_tees:
+                stderr_tees.add(sys.stderr)
+
+        return (
+            cls._get_tee_files(stdout_tees),
+            cls._get_tee_files(stderr_tees),
+            cls._get_tee_files(prompt_tees)
+        )
+
+    def get_pids(self):
+        """ Return the pid for all of the processes in the tree """
+        output = run(
+            ["pstree", "-p", str(self.pid)],
+            stdin=DEVNULL,
+            stdout=PIPE,
+            stderr=DEVNULL,
+            text=True,
+            check=True
+        )
+
+        return re.findall("\\((\\d+)\\)", output.stdout)[::-1]
+
+    def terminate_tree(self):
+        run(
+            ["sudo", "kill"] + self.get_pids(),
+            stdin=DEVNULL, stdout=DEVNULL, stderr=DEVNULL,
+        )
+
+    def kill_tree(self):
+        run(
+            ["sudo", "kill", "-9"] + self.get_pids(),
+            stdin=DEVNULL, stdout=DEVNULL, stderr=DEVNULL,
+        )
+
+    def end_tree(self, sigterm_timeout):
+        """ Try to gracefully terminate the process tree,
+            kill it after 'sigterm_timeout' seconds
+        """
+        if sigterm_timeout:
+            self.terminate_tree()
+            try:
+                self.communicate(timeout=sigterm_timeout)
+            except:  # noqa: E722
+                self.kill_tree()
+        else:
+            self.kill_tree()
+
+    def poll(self):
+        processes = {self.stderr_process, self.stdout_process} - {self}
+        if all(process.poll() is not None for process in processes):
+            return super().poll()
+
+    def wait(self, timeout=None):
+        processes = {self.stderr_process, self.stdout_process} - {self}
+        for process in processes:
+            process.wait(timeout)
+            timeout = None
+        return super().wait(timeout)
+
+    @contextmanager
+    def _streams_restored(self):
+        self.stdout = self.original_stdout
+        self.stderr = self.original_stderr
+        try:
+            yield
+        finally:
+            self.stdout = self.stdout_process.stdout
+            self.stderr = self.stderr_process.stderr
+
+    @contextmanager
+    def _stdin_none(self):
+        original_stdin = self.stdin
+        self.stdin = None
+        try:
+            yield
+        finally:
+            self.stdin = original_stdin
+
+    def _communicate_tees(self, timeout):
+        stdout = stderr = None
+        if self.stdout_process is not self:
+            stdout = self.stdout_process.communicate(timeout=timeout)
+            timeout = None
+        if self.stderr_process is not self:
+            stderr = self.stderr_process.communicate(timeout=timeout)
+            timeout = None
+        return stdout, stderr, timeout
+
+    def _communicate_all(self, timeout):
+        stdout, stderr, timeout = self._communicate_tees(timeout)
+        main_stdout, main_stderr = super().communicate(timeout=timeout)
+        return (
+            main_stdout if self.stdout_process is self else stdout,
+            main_stderr if self.stderr_process is self else stderr
+        )
+
+    @property
+    def _any_communication_started(self):
+        return (
+            self.stdout_process._communication_started or
+            self.stderr_process._communication_started or
+            self._communication_started
+        )
+
+    def communicate(self, input=None, timeout=None):
+        if self._any_communication_started and input:
+            raise ValueError("Cannot send input after starting communication")
+
+        with self._streams_restored():
+            if self.stdin and input:
+                self._stdin_write(input)
+                with self._stdin_none():
+                    return self._communicate_all(timeout)
+            return self._communicate_all(timeout)
+
+
+def normalize_outerr_fds(fds: Iterable[int]):
+    """ Return fds as a set but using 1 and 2 for stdout and stderr file
+        descriptors in case we are being redirected
+    """
+    out_fd = sys.stdout.fileno()  # This might not always be 1
+    err_fd = sys.stderr.fileno()  # This might not always be 2
+    fds = set(fds)
+    if out_fd in fds:
+        fds.remove(out_fd)
+        fds.add(1)
+    if err_fd in fds:
+        fds.remove(err_fd)
+        fds.add(2)
+    return fds
+
+
+def quote(cmd: Iterable[str]):
+    """ Convert the command tokens into a single string that could be pasted into
+        the shell to execute the original command
+    """
+    return " ".join(map(shlex.quote, cmd))
+
+
+def shellify(cmd: Union[str, Iterable[str]], err2out=False, comment=None):
+    """ Quote command if needed and optionally add extra strings to express
+        stderr being redirected to stdout and a comment
+    """
+    if not isinstance(cmd, str):
+        cmd = quote(cmd)
+    if err2out:
+        cmd += " 2>&1"
+    if comment:
+        cmd += f" # {comment}"
+    return cmd
+
+
+# If bash is installed and supports prompt expansion
+if subprocess.run(
+    ["bash", "-c", "echo ${0@P}"],
+    stdin=DEVNULL,
+    stdout=DEVNULL,
+    stderr=DEVNULL
+).returncode == 0:
+    HOME = os.path.expanduser("~")
+    PS1, PS2 = subprocess.run(
+        ["bash", "-ic", "echo \"$PS1\"; echo \"$PS2\""],
+        text=True,
+        stdin=DEVNULL,
+        stdout=PIPE,
+        stderr=DEVNULL
+    ).stdout.splitlines()[-2:]
+
+    def stream_prompts(fds: Iterable[str], cmd, cwd=None, env=None, err2out=False, comment=None):
+        """ Write shell prompt and command into file descriptors fds """
+        fds = normalize_outerr_fds(fds)
+        custom_env = {"CPS1": PS1, "CPS2": PS2}
+        custom_env.update(env or {})
+        custom_env.setdefault("HOME", HOME)
+        script = (
+            "(\n"
+            "    IFS= read -r \"line\"\n"
+            "    echo \"${CPS1@P}${line}\"\n"
+            "    while IFS= read line; do\n"
+            "        echo \"${CPS2@P}${line}\"\n"
+            "    done\n"
+            ") | " + quote(Tee.get_cmd(fds - {1}))
+        )
+        subprocess.run(
+            ["bash", "-c", script],
+            text=True,
+            input=shellify(cmd, err2out, comment),
+            stdout=None if 1 in fds else DEVNULL,
+            stderr=None if 2 in fds else DEVNULL,
+            pass_fds=fds - {1, 2},
+            cwd=cwd,
+            env=custom_env,
+            check=True
+        )
+else:  # Use hard-coded PS1 and PS2 strings
+    def stream_prompts(fds: Iterable[str], cmd, cwd=None, env=None, err2out=False, comment=None):
+        """ Write shell prompt and command into file descriptors fds """
+        cmd = shellify(cmd, err2out, comment) + "\n"
+        input = "$ " + "> ".join(cmd.splitlines(keepends=True))
+        with Tee(PIPE, fds, DEVNULL, text=True) as tee:
+            tee.communicate(input=input)
+
+
+def set_global_echo(value):
+    global GLOBAL_ECHO
+    GLOBAL_ECHO = bool(value)
+
+
+def set_global_stdout_files(*files: TeeStream):
+    global GLOBAL_STDOUTS
+    GLOBAL_STDOUTS = set(*files)
+
+
+def set_global_stderr_files(*files: TeeStream):
+    global GLOBAL_STDERRS
+    GLOBAL_STDERRS = set(*files)
+
+
+def set_global_prompt_files(*files: TeeStream):
+    global GLOBAL_PROMPTS
+    GLOBAL_PROMPTS = set(*files)
+
+
+def _output_context(kwargs, key, encoding, errors, text):
+    """ Create a PipeFile, store it in kwargs[key] and return it if it is FILE.
+        Just return a nullcontext otherwise.
+    """
+    if kwargs.get(key) is FILE:
+        kwargs[key] = PipeFile(encoding=encoding, errors=errors, text=text)
+        return kwargs[key]
+    return nullcontext()
+
+
+def run(
+    args, *, input=None, capture_output=False, timeout=None, check=False,
+    encoding=None, errors=None, text=None, sigterm_timeout=10, **kwargs
+):
+    """ A subprocess.run that instantiates this module's Popen """
+    if input is not None:
+        if kwargs.get("stdin") is not None:
+            raise ValueError("stdin and input arguments may not both be used.")
+        kwargs["stdin"] = PIPE
+
+    if capture_output:
+        if kwargs.get("stdout") is not None or kwargs.get("stderr") is not None:
+            raise ValueError("stdout and stderr arguments may not be used with capture_output.")
+        kwargs["stdout"] = FILE
+        kwargs["stderr"] = FILE
+
+    comment = f"timeout={timeout}" if timeout else None
+    with (
+        _output_context(kwargs, "stdout", encoding, errors, text) as stdout_file,
+        _output_context(kwargs, "stderr", encoding, errors, text) as stderr_file,
+        Popen(args, encoding=encoding, errors=errors, text=text,
+              comment=comment, **kwargs) as process
+    ):
+        start = psutil.Process(process.pid).create_time()
+        try:
+            stdout, stderr = process.communicate(input, timeout=timeout)
+        except TimeoutExpired:
+            process.end_tree(sigterm_timeout)
+            process.wait()
+            raise
+        except:  # noqa: E722
+            process.kill_tree()
+            raise
+        finally:
+            end = time()
+            returncode = process.poll()
+
+        if stdout_file:
+            stdout = stdout_file.read_all()
+        if stderr_file:
+            stderr = stderr_file.read_all()
+
+        if check and returncode:
+            raise CalledProcessError(returncode, process.args, output=stdout, stderr=stderr)
+
+        output = CompletedProcess(process.args, returncode, stdout, stderr)
+        output.time = end - start
+        return output
+
+
+def call(*args, **kwargs):
+    return run(*args, **kwargs).returncode
+
+
+def check_call(*args, **kwargs):
+    kwargs["check"] = True
+    return call(*args, **kwargs)
+
+
+def check_output(*args, **kwargs):
+    kwargs["check"] = True
+    kwargs.setdefault("stdout", PIPE)
+    return run(*args, **kwargs).stdout
+
+
+def getoutput(*args, **kwargs):
+    return getstatusoutput(*args, **kwargs)[1]
+
+
+def getstatusoutput(*args, **kwargs):
+    kwargs["stdout"] = PIPE
+    kwargs["stderr"] = STDOUT
+    kwargs["shell"] = True
+    kwargs["text"] = True
+    output = run(*args, **kwargs)
+    return (output.returncode, output.stdout)