1 

     2 """

     3     ForkedFunc provides a way to run a function in a forked process

     4     and get at its return value, stdout and stderr output as well

     5     as signals and exitstatusus.

     6 

     7     XXX see if tempdir handling is sane

     8 """

     9 

    10 import py

    11 import os

    12 import sys

    13 import marshal

    14 

    15 class ForkedFunc(object):

    16     EXITSTATUS_EXCEPTION = 3

    17     def __init__(self, fun, args=None, kwargs=None, nice_level=0):

    18         if args is None:

    19             args = []

    20         if kwargs is None:

    21             kwargs = {}

    22         self.fun = fun

    23         self.args = args

    24         self.kwargs = kwargs

    25         self.tempdir = tempdir = py.path.local.mkdtemp()

    26         self.RETVAL = tempdir.ensure('retval')

    27         self.STDOUT = tempdir.ensure('stdout')

    28         self.STDERR = tempdir.ensure('stderr')

    29 

    30         pid = os.fork()

    31         if pid: # in parent process

    32             self.pid = pid

    33         else: # in child process

    34             self._child(nice_level)

    35 

    36     def _child(self, nice_level):

    37         # right now we need to call a function, but first we need to

    38         # map all IO that might happen

    39         # make sure sys.stdout points to file descriptor one

    40         sys.stdout = stdout = self.STDOUT.open('w')

    41         sys.stdout.flush()

    42         fdstdout = stdout.fileno()

    43         if fdstdout != 1:

    44             os.dup2(fdstdout, 1)

    45         sys.stderr = stderr = self.STDERR.open('w')

    46         fdstderr = stderr.fileno()

    47         if fdstderr != 2:

    48             os.dup2(fdstderr, 2)

    49         retvalf = self.RETVAL.open("wb")

    50         EXITSTATUS = 0

    51         try:

    52             if nice_level:

    53                 os.nice(nice_level)

    54             try:

    55                 retval = self.fun(*self.args, **self.kwargs)

    56                 retvalf.write(marshal.dumps(retval))

    57             except:

    58                 excinfo = py.code.ExceptionInfo()

    59                 stderr.write(excinfo.exconly())

    60                 EXITSTATUS = self.EXITSTATUS_EXCEPTION

    61         finally:

    62             stdout.close()

    63             stderr.close()

    64             retvalf.close()

    65         os.close(1)

    66         os.close(2)

    67         os._exit(EXITSTATUS)

    68 

    69     def waitfinish(self, waiter=os.waitpid):

    70         pid, systemstatus = waiter(self.pid, 0)

    71         if systemstatus:

    72             if os.WIFSIGNALED(systemstatus):

    73                 exitstatus = os.WTERMSIG(systemstatus) + 128

    74             else:

    75                 exitstatus = os.WEXITSTATUS(systemstatus)

    76             #raise ExecutionFailed(status, systemstatus, cmd,

    77             #                      ''.join(out), ''.join(err))

    78         else:

    79             exitstatus = 0

    80         signal = systemstatus & 0x7f

    81         if not exitstatus and not signal:

    82             retval = self.RETVAL.open('rb')

    83             try:

    84                 retval_data = retval.read()

    85             finally:

    86                 retval.close()

    87             retval = marshal.loads(retval_data)

    88         else:

    89             retval = None

    90         stdout = self.STDOUT.read()

    91         stderr = self.STDERR.read()

    92         self._removetemp()

    93         return Result(exitstatus, signal, retval, stdout, stderr)

    94 

    95     def _removetemp(self):

    96         if self.tempdir.check():

    97             self.tempdir.remove()

    98 

    99     def __del__(self):

   100         self._removetemp()

   101 

   102 class Result(object):

   103     def __init__(self, exitstatus, signal, retval, stdout, stderr):

   104         self.exitstatus = exitstatus

   105         self.signal = signal

   106         self.retval = retval

   107         self.out = stdout

   108         self.err = stderr