| |
| """ |
| 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 |
