twisted/internet/utils.py

   1 # -*- test-case-name: twisted.test.test_iutils -*-
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5 """
   6 Utility methods.
   7 """
   8
   9 import sys, warnings
  10
  11 from twisted.internet import protocol, defer
  12 from twisted.python import failure, util as tputil
  13
  14 try:
  15     import cStringIO as StringIO
  16 except ImportError:
  17     import StringIO
  18
  19 def _callProtocolWithDeferred(protocol, executable, args, env, path, reactor=None):
  20     if reactor is None:
  21         from twisted.internet import reactor
  22
  23     d = defer.Deferred()
  24     p = protocol(d)
  25     reactor.spawnProcess(p, executable, (executable,)+tuple(args), env, path)
  26     return d
  27
  28
  29
  30 class _UnexpectedErrorOutput(IOError):
  31     """
  32     Standard error data was received where it was not expected.  This is a
  33     subclass of L{IOError} to preserve backward compatibility with the previous
  34     error behavior of L{getProcessOutput}.
  35
  36     @ivar processEnded: A L{Deferred} which will fire when the process which
  37         produced the data on stderr has ended (exited and all file descriptors
  38         closed).
  39     """
  40     def __init__(self, text, processEnded):
  41         IOError.__init__(self, "got stderr: %r" % (text,))
  42         self.processEnded = processEnded
  43
  44
  45
  46 class _BackRelay(protocol.ProcessProtocol):
  47     """
  48     Trivial protocol for communicating with a process and turning its output
  49     into the result of a L{Deferred}.
  50
  51     @ivar deferred: A L{Deferred} which will be called back with all of stdout
  52         and, if C{errortoo} is true, all of stderr as well (mixed together in
  53         one string).  If C{errortoo} is false and any bytes are received over
  54         stderr, this will fire with an L{_UnexpectedErrorOutput} instance and
  55         the attribute will be set to C{None}.
  56
  57     @ivar onProcessEnded: If C{errortoo} is false and bytes are received over
  58         stderr, this attribute will refer to a L{Deferred} which will be called
  59         back when the process ends.  This C{Deferred} is also associated with
  60         the L{_UnexpectedErrorOutput} which C{deferred} fires with earlier in
  61         this case so that users can determine when the process has actually
  62         ended, in addition to knowing when bytes have been received via stderr.
  63     """
  64
  65     def __init__(self, deferred, errortoo=0):
  66         self.deferred = deferred
  67         self.s = StringIO.StringIO()
  68         if errortoo:
  69             self.errReceived = self.errReceivedIsGood
  70         else:
  71             self.errReceived = self.errReceivedIsBad
  72
  73     def errReceivedIsBad(self, text):
  74         if self.deferred is not None:
  75             self.onProcessEnded = defer.Deferred()
  76             err = _UnexpectedErrorOutput(text, self.onProcessEnded)
  77             self.deferred.errback(failure.Failure(err))
  78             self.deferred = None
  79             self.transport.loseConnection()
  80
  81     def errReceivedIsGood(self, text):
  82         self.s.write(text)
  83
  84     def outReceived(self, text):
  85         self.s.write(text)
  86
  87     def processEnded(self, reason):
  88         if self.deferred is not None:
  89             self.deferred.callback(self.s.getvalue())
  90         elif self.onProcessEnded is not None:
  91             self.onProcessEnded.errback(reason)
  92
  93
  94
  95 def getProcessOutput(executable, args=(), env={}, path=None, reactor=None,
  96                      errortoo=0):
  97     """
  98     Spawn a process and return its output as a deferred returning a string.
  99
 100     @param executable: The file name to run and get the output of - the
 101                        full path should be used.
 102
 103     @param args: the command line arguments to pass to the process; a
 104                  sequence of strings. The first string should *NOT* be the
 105                  executable's name.
 106
 107     @param env: the environment variables to pass to the processs; a
 108                 dictionary of strings.
 109
 110     @param path: the path to run the subprocess in - defaults to the
 111                  current directory.
 112
 113     @param reactor: the reactor to use - defaults to the default reactor
 114
 115     @param errortoo: If true, include stderr in the result.  If false, if
 116         stderr is received the returned L{Deferred} will errback with an
 117         L{IOError} instance with a C{processEnded} attribute.  The
 118         C{processEnded} attribute refers to a L{Deferred} which fires when the
 119         executed process ends.
 120     """
 121     return _callProtocolWithDeferred(lambda d:
 122                                         _BackRelay(d, errortoo=errortoo),
 123                                      executable, args, env, path,
 124                                      reactor)
 125
 126
 127 class _ValueGetter(protocol.ProcessProtocol):
 128
 129     def __init__(self, deferred):
 130         self.deferred = deferred
 131
 132     def processEnded(self, reason):
 133         self.deferred.callback(reason.value.exitCode)
 134
 135
 136 def getProcessValue(executable, args=(), env={}, path=None, reactor=None):
 137     """Spawn a process and return its exit code as a Deferred."""
 138     return _callProtocolWithDeferred(_ValueGetter, executable, args, env, path,
 139                                     reactor)
 140
 141
 142 class _EverythingGetter(protocol.ProcessProtocol):
 143
 144     def __init__(self, deferred):
 145         self.deferred = deferred
 146         self.outBuf = StringIO.StringIO()
 147         self.errBuf = StringIO.StringIO()
 148         self.outReceived = self.outBuf.write
 149         self.errReceived = self.errBuf.write
 150
 151     def processEnded(self, reason):
 152         out = self.outBuf.getvalue()
 153         err = self.errBuf.getvalue()
 154         e = reason.value
 155         code = e.exitCode
 156         if e.signal:
 157             self.deferred.errback((out, err, e.signal))
 158         else:
 159             self.deferred.callback((out, err, code))
 160
 161 def getProcessOutputAndValue(executable, args=(), env={}, path=None,
 162                              reactor=None):
 163     """Spawn a process and returns a Deferred that will be called back with
 164     its output (from stdout and stderr) and it's exit code as (out, err, code)
 165     If a signal is raised, the Deferred will errback with the stdout and
 166     stderr up to that point, along with the signal, as (out, err, signalNum)
 167     """
 168     return _callProtocolWithDeferred(_EverythingGetter, executable, args, env, path,
 169                                     reactor)
 170
 171 def _resetWarningFilters(passthrough, addedFilters):
 172     for f in addedFilters:
 173         try:
 174             warnings.filters.remove(f)
 175         except ValueError:
 176             pass
 177     return passthrough
 178
 179
 180 def runWithWarningsSuppressed(suppressedWarnings, f, *a, **kw):
 181     """Run the function C{f}, but with some warnings suppressed.
 182
 183     @param suppressedWarnings: A list of arguments to pass to filterwarnings.
 184                                Must be a sequence of 2-tuples (args, kwargs).
 185     @param f: A callable, followed by its arguments and keyword arguments
 186     """
 187     for args, kwargs in suppressedWarnings:
 188         warnings.filterwarnings(*args, **kwargs)
 189     addedFilters = warnings.filters[:len(suppressedWarnings)]
 190     try:
 191         result = f(*a, **kw)
 192     except:
 193         exc_info = sys.exc_info()
 194         _resetWarningFilters(None, addedFilters)
 195         raise exc_info[0], exc_info[1], exc_info[2]
 196     else:
 197         if isinstance(result, defer.Deferred):
 198             result.addBoth(_resetWarningFilters, addedFilters)
 199         else:
 200             _resetWarningFilters(None, addedFilters)
 201         return result
 202
 203
 204 def suppressWarnings(f, *suppressedWarnings):
 205     """
 206     Wrap C{f} in a callable which suppresses the indicated warnings before
 207     invoking C{f} and unsuppresses them afterwards.  If f returns a Deferred,
 208     warnings will remain suppressed until the Deferred fires.
 209     """
 210     def warningSuppressingWrapper(*a, **kw):
 211         return runWithWarningsSuppressed(suppressedWarnings, f, *a, **kw)
 212     return tputil.mergeFunctionMetadata(f, warningSuppressingWrapper)
 213
 214
 215 __all__ = [
 216     "runWithWarningsSuppressed", "suppressWarnings",
 217
 218     "getProcessOutput", "getProcessValue", "getProcessOutputAndValue",
 219     ]