109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
|
|
"""
|
|
ForkedFunc provides a way to run a function in a forked process
|
|
and get at its return value, stdout and stderr output as well
|
|
as signals and exitstatusus.
|
|
|
|
XXX see if tempdir handling is sane
|
|
"""
|
|
|
|
import py
|
|
import os
|
|
import sys
|
|
import marshal
|
|
|
|
class ForkedFunc(object):
|
|
EXITSTATUS_EXCEPTION = 3
|
|
def __init__(self, fun, args=None, kwargs=None, nice_level=0):
|
|
if args is None:
|
|
args = []
|
|
if kwargs is None:
|
|
kwargs = {}
|
|
self.fun = fun
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
self.tempdir = tempdir = py.path.local.mkdtemp()
|
|
self.RETVAL = tempdir.ensure('retval')
|
|
self.STDOUT = tempdir.ensure('stdout')
|
|
self.STDERR = tempdir.ensure('stderr')
|
|
|
|
pid = os.fork()
|
|
if pid: # in parent process
|
|
self.pid = pid
|
|
else: # in child process
|
|
self._child(nice_level)
|
|
|
|
def _child(self, nice_level):
|
|
# right now we need to call a function, but first we need to
|
|
# map all IO that might happen
|
|
# make sure sys.stdout points to file descriptor one
|
|
sys.stdout = stdout = self.STDOUT.open('w')
|
|
sys.stdout.flush()
|
|
fdstdout = stdout.fileno()
|
|
if fdstdout != 1:
|
|
os.dup2(fdstdout, 1)
|
|
sys.stderr = stderr = self.STDERR.open('w')
|
|
fdstderr = stderr.fileno()
|
|
if fdstderr != 2:
|
|
os.dup2(fdstderr, 2)
|
|
retvalf = self.RETVAL.open("wb")
|
|
EXITSTATUS = 0
|
|
try:
|
|
if nice_level:
|
|
os.nice(nice_level)
|
|
try:
|
|
retval = self.fun(*self.args, **self.kwargs)
|
|
retvalf.write(marshal.dumps(retval))
|
|
except:
|
|
excinfo = py.code.ExceptionInfo()
|
|
stderr.write(excinfo.exconly())
|
|
EXITSTATUS = self.EXITSTATUS_EXCEPTION
|
|
finally:
|
|
stdout.close()
|
|
stderr.close()
|
|
retvalf.close()
|
|
os.close(1)
|
|
os.close(2)
|
|
os._exit(EXITSTATUS)
|
|
|
|
def waitfinish(self, waiter=os.waitpid):
|
|
pid, systemstatus = waiter(self.pid, 0)
|
|
if systemstatus:
|
|
if os.WIFSIGNALED(systemstatus):
|
|
exitstatus = os.WTERMSIG(systemstatus) + 128
|
|
else:
|
|
exitstatus = os.WEXITSTATUS(systemstatus)
|
|
#raise ExecutionFailed(status, systemstatus, cmd,
|
|
# ''.join(out), ''.join(err))
|
|
else:
|
|
exitstatus = 0
|
|
signal = systemstatus & 0x7f
|
|
if not exitstatus and not signal:
|
|
retval = self.RETVAL.open('rb')
|
|
try:
|
|
retval_data = retval.read()
|
|
finally:
|
|
retval.close()
|
|
retval = marshal.loads(retval_data)
|
|
else:
|
|
retval = None
|
|
stdout = self.STDOUT.read()
|
|
stderr = self.STDERR.read()
|
|
self._removetemp()
|
|
return Result(exitstatus, signal, retval, stdout, stderr)
|
|
|
|
def _removetemp(self):
|
|
if self.tempdir.check():
|
|
self.tempdir.remove()
|
|
|
|
def __del__(self):
|
|
self._removetemp()
|
|
|
|
class Result(object):
|
|
def __init__(self, exitstatus, signal, retval, stdout, stderr):
|
|
self.exitstatus = exitstatus
|
|
self.signal = signal
|
|
self.retval = retval
|
|
self.out = stdout
|
|
self.err = stderr
|