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