"""
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.
"""
import py
import os
import sys
import marshal
def get_unbuffered_io(fd, filename):
f = open(str(filename), "w")
if fd != f.fileno():
os.dup2(f.fileno(), fd)
class AutoFlush:
def write(self, data):
f.write(data)
f.flush()
def __getattr__(self, name):
return getattr(f, name)
return AutoFlush()
class ForkedFunc:
EXITSTATUS_EXCEPTION = 3
def __init__(self, fun, args=None, kwargs=None, nice_level=0,
child_on_start=None, child_on_exit=None):
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.pid = None
self._child(nice_level, child_on_start, child_on_exit)
def _child(self, nice_level, child_on_start, child_on_exit):
# right now we need to call a function, but first we need to
# map all IO that might happen
sys.stdout = stdout = get_unbuffered_io(1, self.STDOUT)
sys.stderr = stderr = get_unbuffered_io(2, self.STDERR)
retvalf = self.RETVAL.open("wb")
EXITSTATUS = 0
try:
if nice_level:
os.nice(nice_level)
try:
if child_on_start is not None:
child_on_start()
retval = self.fun(*self.args, **self.kwargs)
retvalf.write(marshal.dumps(retval))
if child_on_exit is not None:
child_on_exit()
except:
excinfo = py.code.ExceptionInfo()
stderr.write(str(excinfo._getreprcrash()))
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)
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):
if self.pid is not None: # only clean up in main process
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
