summaryrefslogtreecommitdiff
path: root/exec_utils.py
blob: 89cfbd7bf15f23a7a0effaed24954f7b4480aded (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
#!/usr/bin/env python3

import atexit
import logging
import shlex
import subprocess
from typing import List, Optional


logger = logging.getLogger(__file__)


def cmd_with_timeout(command: str, timeout_seconds: Optional[float]) -> int:
    """
    Run a command but do not let it run for more than timeout seconds.

    >>> cmd_with_timeout('/bin/echo foo', 10.0)
    0

    >>> cmd_with_timeout('/bin/sleep 2', 0.1)
    Traceback (most recent call last):
    ...
    subprocess.TimeoutExpired: Command '['/bin/bash', '-c', '/bin/sleep 2']' timed out after 0.1 seconds

    """
    return subprocess.check_call(
        ["/bin/bash", "-c", command], timeout=timeout_seconds
    )


def cmd(command: str, timeout_seconds: Optional[float] = None) -> str:
    """Run a command with everything encased in a string and return
    the output text as a string.  Raises subprocess.CalledProcessError.

    >>> cmd('/bin/echo foo')[:-1]
    'foo'

    >>> cmd('/bin/sleep 2', 0.1)
    Traceback (most recent call last):
    ...
    subprocess.TimeoutExpired: Command '/bin/sleep 2' timed out after 0.1 seconds

    """
    ret = subprocess.run(
        command, shell=True, capture_output=True, check=True, timeout=timeout_seconds,
    ).stdout
    return ret.decode("utf-8")


def run_silently(command: str) -> None:
    """Run a command silently but raise subprocess.CalledProcessError if
    it fails.

    >>> run_silently("/usr/bin/true")

    >>> run_silently("/usr/bin/false")
    Traceback (most recent call last):
    ...
    subprocess.CalledProcessError: Command '/usr/bin/false' returned non-zero exit status 1.

    """
    subprocess.run(
        command, shell=True, stderr=subprocess.DEVNULL,
        stdout=subprocess.DEVNULL, capture_output=False, check=True
    )


def cmd_in_background(
        command: str, *, silent: bool = False
) -> subprocess.Popen:
    args = shlex.split(command)
    if silent:
        subproc = subprocess.Popen(args,
                                   stdin=subprocess.DEVNULL,
                                   stdout=subprocess.DEVNULL,
                                   stderr=subprocess.DEVNULL)
    else:
        subproc = subprocess.Popen(args, stdin=subprocess.DEVNULL)

    def kill_subproc() -> None:
        try:
            if subproc.poll() is None:
                logger.info("At exit handler: killing {}: {}".format(subproc, command))
                subproc.terminate()
                subproc.wait(timeout=10.0)
        except BaseException as be:
            logger.exception(be)
    atexit.register(kill_subproc)
    return subproc


def cmd_list(command: List[str]) -> str:
    """Run a command with args encapsulated in a list and return the
    output text as a string.  Raises subprocess.CalledProcessError.
    """
    ret = subprocess.run(command, capture_output=True, check=True).stdout
    return ret.decode("utf-8")


if __name__ == '__main__':
    import doctest
    doctest.testmod()