shell.py 4.46 KB
Newer Older
Johannes Bechberger's avatar
Johannes Bechberger committed
1 2 3 4 5 6 7
# source: https://github.com/libfirm/sisyphus

"""
Convenience function
Alternative to subprocess and os.system
"""
import subprocess
8 9 10 11 12 13 14
has_resource_module = True
try:
	import resource
except ImportError:
	has_resource_module = False
	pass
	
Johannes Bechberger's avatar
Johannes Bechberger committed
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
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):
33 34
    if not has_resource_module:
        return
Johannes Bechberger's avatar
Johannes Bechberger committed
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
    (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):
62 63
        if not has_resource_module:
            return
Johannes Bechberger's avatar
Johannes Bechberger committed
64 65 66 67 68 69 70
        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:
71 72 73
            prexec_args = {}
            if has_resource_module:
                prexec_args["preexec_fn"] = self._set_rlimit
Johannes Bechberger's avatar
Johannes Bechberger committed
74 75 76
            self.proc = subprocess.Popen(self.cmd,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE,
77 78
                                         env=self.env,
                                         **prexec_args)
Johannes Bechberger's avatar
Johannes Bechberger committed
79 80 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
            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:
        cmd = filter(lambda x: x, cmd.split(' '))
123 124
    else:
        cmd = list(filter(lambda x: x, cmd[0].split(' '))) + cmd[1:]
Johannes Bechberger's avatar
Johannes Bechberger committed
125 126 127 128 129 130 131 132 133 134
    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)