from twisted.internet.defer import DeferredList, Deferred
from twisted.internet.protocol import ProcessProtocol
from twisted.internet.process import Process, ProcessReader
from twisted.internet import reactor

import cStringIO


def gather_results(deferreds, consume_errors=False):
d = DeferredList(deferreds, fireOnOneErrback=1,
consumeErrors=consume_errors)
d.addCallback(lambda r: [x[1] for x in r])
d.addErrback(lambda f: f.value.subFailure)
return d


class AllOutputProcessProtocol(ProcessProtocol):
"""A process protocol for getting stdout, stderr and exit code."""

def __init__(self, deferred, stdin=None, line_received=None):
self.deferred = deferred
self.outBuf = cStringIO.StringIO()
self.errBuf = cStringIO.StringIO()
self.errReceived = self.errBuf.write
self.stdin = stdin
self.line_received = line_received
self._partial_line = ""

def connectionMade(self):
if self.stdin is not None:
self.transport.write(self.stdin)
self.transport.closeStdin()

def outReceived(self, data):
self.outBuf.write(data)

if self.line_received is None:
return

# data may contain more than one line, so we split the output and save
# the last line. If it's an empty string nothing happens, otherwise it
# will be returned once complete
lines = data.split("\n")
lines[0] = self._partial_line + lines[0]
self._partial_line = lines.pop()

for line in lines:
self.line_received(line)

def processEnded(self, reason):
if self._partial_line:
self.line_received(self._partial_line)
self._partial_line = ""
out = self.outBuf.getvalue()
err = self.errBuf.getvalue()
e = reason.value
code = e.exitCode
if e.signal:
self.deferred.errback((out, err, e.signal))
else:
self.deferred.callback((out, err, code))


def spawn_process(executable, args=(), env={}, path=None, uid=None, gid=None,
usePTY=False, wait_pipes=True, line_received=None,
stdin=None):
"""
Spawn a process using Twisted reactor.

Return a deferred which will be called with process stdout, stderr and exit
code.

@param wait_pipes: if set to False, don't wait for stdin/stdout pipes to
close when process ends.
@param line_received: an optional callback called with every line of
output from the process as parameter.

@note: compared to reactor.spawnProcess, this version does NOT require the
executable name as first element of args.
"""

list_args = [executable]
list_args.extend(args)

result = Deferred()
protocol = AllOutputProcessProtocol(result, stdin=stdin,
line_received=line_received)
process = reactor.spawnProcess(protocol, executable, args=list_args,
env=env, path=path, uid=uid, gid=gid,
usePTY=usePTY)

if not wait_pipes:

def maybeCallProcessEnded():
"""A less strict version of Process.maybeCallProcessEnded.

This behaves exactly like the original method, but in case the
process has ended already and sent us a SIGCHLD, it doesn't wait
for the stdin/stdout pipes to close, because the child process
itself might have passed them to its own child processes.

@note: Twisted 8.2 now has a processExited hook that could
be used in place of this workaround.
"""
if process.pipes and not process.pid:
for pipe in process.pipes.itervalues():
if isinstance(pipe, ProcessReader):
# Read whatever is left
pipe.doRead()
pipe.stopReading()
process.pipes = {}
Process.maybeCallProcessEnded(process)

process.maybeCallProcessEnded = maybeCallProcessEnded

return result