shell.py 4.63 KB
Newer Older
Johannes Bechberger's avatar
Johannes Bechberger committed
1
# source: https://github.com/libfirm/sisyphus, but slightly modified to support Windows
Johannes Bechberger's avatar
Johannes Bechberger committed
2
3
4
5
6

"""
Convenience function
Alternative to subprocess and os.system
"""
Johannes Bechberger's avatar
Johannes Bechberger committed
7
import shlex
Johannes Bechberger's avatar
Johannes Bechberger committed
8
import subprocess
9
10
has_resource_module = True
try:
Johannes Bechberger's avatar
Johannes Bechberger committed
11
    import resource
12
except ImportError:
Johannes Bechberger's avatar
Johannes Bechberger committed
13
14
    has_resource_module = False
    pass
15
	
Johannes Bechberger's avatar
Johannes Bechberger committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import sys
import signal
import threading

import logging
_LOG = logging.getLogger("sisyphus")

_EXIT_CODES = dict((-k, v) for v, k in signal.__dict__.items() if v.startswith('SIG'))
del _EXIT_CODES[0]


class SigKill(Exception):
    def __init__(self, retcode, name):
        self.retcode = retcode
        self.name = name


def _lower_rlimit(res, limit):
34
35
    if not has_resource_module:
        return
Johannes Bechberger's avatar
Johannes Bechberger committed
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
    (soft, hard) = resource.getrlimit(res)
    if soft > limit or soft == resource.RLIM_INFINITY:
        soft = limit
    if hard > limit or hard == resource.RLIM_INFINITY:
        hard = limit
    resource.setrlimit(res, (soft, hard))


class _Execute(object):
    def __init__(self, cmd, timeout, env, rlimit):
        self.cmd       = cmd
        self.timeout   = timeout
        self.env       = env
        self.proc      = None
        self.exception = None
        self.rlimit    = rlimit
        MB = 1024 * 1024
        if not 'RLIMIT_CORE' in rlimit:
            rlimit['RLIMIT_CORE'] = 0
        if not 'RLIMIT_DATA' in rlimit:
            rlimit['RLIMIT_DATA'] = 1024 * MB
        if not 'RLIMIT_STACK' in rlimit:
            rlimit['RLIMIT_STACK'] = 1024 * MB
        if not 'RLIMIT_FSIZE' in rlimit:
            rlimit['RLIMIT_FSIZE'] = 512 * MB

    def _set_rlimit(self):
63
64
        if not has_resource_module:
            return
Johannes Bechberger's avatar
Johannes Bechberger committed
65
66
67
68
69
70
71
        if self.timeout > 0.0:
            _lower_rlimit(resource.RLIMIT_CPU, self.timeout)
        for k,v in self.rlimit.items():
            _lower_rlimit(getattr(resource,k), v)

    def _run_process(self):
        try:
72
73
74
            prexec_args = {}
            if has_resource_module:
                prexec_args["preexec_fn"] = self._set_rlimit
Johannes Bechberger's avatar
Johannes Bechberger committed
75
76
77
            self.proc = subprocess.Popen(self.cmd,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE,
78
                                         env=self.env,
Johannes Bechberger's avatar
Johannes Bechberger committed
79
                                         shell=True,
80
                                         **prexec_args)
Johannes Bechberger's avatar
Johannes Bechberger committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
            x = self.proc.communicate()
            self.out, self.err = x
            self.returncode = self.proc.returncode
        except Exception as e:
            self.exception = e

    def run(self):
        if self.timeout > 0.0:
            _LOG.debug("run with timeout %.1f" % float(self.timeout))
            thread = threading.Thread(target=self._run_process)
            thread.start()
            thread.join(float(self.timeout))
            if self.exception is not None:
                raise self.exception
            if thread.is_alive():
                _LOG.debug("timeout %.1f reached, terminate process" % self.timeout)
                if self.proc is not None:
                    self.proc.terminate()
                    thread.join(1.0)
                    if thread.is_alive():
                        _LOG.debug("termination ignored, now kill process")
                        self.proc.kill()
                        thread.join()
                raise SigKill(signal.SIGXCPU, "SIGXCPU")
        else:
            self._run_process()

        # Usually python can recognize application terminations triggered by
        # signals, but somehow it doesn't do this for java (I suspect, that java
        # catches the signals itself but then shuts down more 'cleanly' than it
        # should. Calculate to python convention returncode
        if self.returncode > 127:
            self.returncode = 128 - self.returncode
        if self.returncode in _EXIT_CODES:
            _LOG.debug("returncode %d recognized as '%s'" % (self.returncode, _EXIT_CODES[self.returncode]))
            raise SigKill(self.returncode, _EXIT_CODES[self.returncode])
        return (self.out, self.err, self.returncode)

def execute(cmd, env=None, timeout=0, rlimit=None, propagate_sigint=True):
    """Execute a command and return stderr and stdout data"""
    if not rlimit:
        rlimit = dict()
    if cmd is str:
Johannes Bechberger's avatar
Johannes Bechberger committed
124
125
        pass
        #cmd = filter(lambda x: x, cmd.split(' '))
126
    else:
Johannes Bechberger's avatar
Johannes Bechberger committed
127
128
        #cmd = shlex.split(cmd[0]) + cmd[1:]
        cmd = cmd[0] + " " + " ".join(shlex.quote(c) for c in cmd[1:])
Johannes Bechberger's avatar
Johannes Bechberger committed
129
130
131
132
133
134
135
136
137
138
    exc = _Execute(cmd, timeout, env, rlimit)
    (out, err, returncode) = exc.run()
    if returncode == -signal.SIGINT:
        raise KeyboardInterrupt
    return (out, err, returncode)


if __name__ == "__main__":
    out, err, retcode = execute("hostname")
    print (out)