subprocess.py 7.45 KB
Newer Older
1 2 3 4 5 6 7 8
__all__ = ['create_subprocess_exec', 'create_subprocess_shell']

import subprocess

from . import events
from . import protocols
from . import streams
from . import tasks
9
from .coroutines import coroutine
10
from .log import logger
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26


PIPE = subprocess.PIPE
STDOUT = subprocess.STDOUT
DEVNULL = subprocess.DEVNULL


class SubprocessStreamProtocol(streams.FlowControlMixin,
                               protocols.SubprocessProtocol):
    """Like StreamReaderProtocol, but for a subprocess."""

    def __init__(self, limit, loop):
        super().__init__(loop=loop)
        self._limit = limit
        self.stdin = self.stdout = self.stderr = None
        self._transport = None
27 28
        self._process_exited = False
        self._pipe_fds = []
29

30 31 32 33 34 35 36 37 38 39
    def __repr__(self):
        info = [self.__class__.__name__]
        if self.stdin is not None:
            info.append('stdin=%r' % self.stdin)
        if self.stdout is not None:
            info.append('stdout=%r' % self.stdout)
        if self.stderr is not None:
            info.append('stderr=%r' % self.stderr)
        return '<%s>' % ' '.join(info)

40 41
    def connection_made(self, transport):
        self._transport = transport
42 43 44

        stdout_transport = transport.get_pipe_transport(1)
        if stdout_transport is not None:
45 46
            self.stdout = streams.StreamReader(limit=self._limit,
                                               loop=self._loop)
47
            self.stdout.set_transport(stdout_transport)
48
            self._pipe_fds.append(1)
49 50 51

        stderr_transport = transport.get_pipe_transport(2)
        if stderr_transport is not None:
52 53
            self.stderr = streams.StreamReader(limit=self._limit,
                                               loop=self._loop)
54
            self.stderr.set_transport(stderr_transport)
55
            self._pipe_fds.append(2)
56 57 58 59

        stdin_transport = transport.get_pipe_transport(0)
        if stdin_transport is not None:
            self.stdin = streams.StreamWriter(stdin_transport,
60 61 62
                                              protocol=self,
                                              reader=None,
                                              loop=self._loop)
63

64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
    def pipe_data_received(self, fd, data):
        if fd == 1:
            reader = self.stdout
        elif fd == 2:
            reader = self.stderr
        else:
            reader = None
        if reader is not None:
            reader.feed_data(data)

    def pipe_connection_lost(self, fd, exc):
        if fd == 0:
            pipe = self.stdin
            if pipe is not None:
                pipe.close()
            self.connection_lost(exc)
            return
        if fd == 1:
            reader = self.stdout
        elif fd == 2:
            reader = self.stderr
        else:
            reader = None
        if reader != None:
            if exc is None:
                reader.feed_eof()
            else:
                reader.set_exception(exc)
92

93 94 95
        if fd in self._pipe_fds:
            self._pipe_fds.remove(fd)
        self._maybe_close_transport()
96 97

    def process_exited(self):
98 99
        self._process_exited = True
        self._maybe_close_transport()
100

101 102 103 104
    def _maybe_close_transport(self):
        if len(self._pipe_fds) == 0 and self._process_exited:
            self._transport.close()
            self._transport = None
105

106 107 108 109 110 111 112 113 114 115 116

class Process:
    def __init__(self, transport, protocol, loop):
        self._transport = transport
        self._protocol = protocol
        self._loop = loop
        self.stdin = protocol.stdin
        self.stdout = protocol.stdout
        self.stderr = protocol.stderr
        self.pid = transport.get_pid()

117 118 119
    def __repr__(self):
        return '<%s %s>' % (self.__class__.__name__, self.pid)

120 121 122 123
    @property
    def returncode(self):
        return self._transport.get_returncode()

124
    @coroutine
125
    def wait(self):
126
        """Wait until the process exit and return the process return code.
127

128
        This method is a coroutine."""
129
        return (yield from self._transport._wait())
130 131 132 133 134 135 136 137 138 139

    def send_signal(self, signal):
        self._transport.send_signal(signal)

    def terminate(self):
        self._transport.terminate()

    def kill(self):
        self._transport.kill()

140
    @coroutine
141
    def _feed_stdin(self, input):
142
        debug = self._loop.get_debug()
143
        self.stdin.write(input)
144
        if debug:
145 146
            logger.debug('%r communicate: feed stdin (%s bytes)',
                        self, len(input))
147 148
        try:
            yield from self.stdin.drain()
149 150 151 152
        except (BrokenPipeError, ConnectionResetError) as exc:
            # communicate() ignores BrokenPipeError and ConnectionResetError
            if debug:
                logger.debug('%r communicate: stdin got %r', self, exc)
153

154
        if debug:
155
            logger.debug('%r communicate: close stdin', self)
156 157
        self.stdin.close()

158
    @coroutine
159 160 161
    def _noop(self):
        return None

162
    @coroutine
163 164 165 166 167 168 169
    def _read_stream(self, fd):
        transport = self._transport.get_pipe_transport(fd)
        if fd == 2:
            stream = self.stderr
        else:
            assert fd == 1
            stream = self.stdout
170 171 172
        if self._loop.get_debug():
            name = 'stdout' if fd == 1 else 'stderr'
            logger.debug('%r communicate: read %s', self, name)
173
        output = yield from stream.read()
174 175 176
        if self._loop.get_debug():
            name = 'stdout' if fd == 1 else 'stderr'
            logger.debug('%r communicate: close %s', self, name)
177 178 179
        transport.close()
        return output

180
    @coroutine
181
    def communicate(self, input=None):
182
        if input is not None:
183 184 185 186 187 188 189 190 191 192 193 194
            stdin = self._feed_stdin(input)
        else:
            stdin = self._noop()
        if self.stdout is not None:
            stdout = self._read_stream(1)
        else:
            stdout = self._noop()
        if self.stderr is not None:
            stderr = self._read_stream(2)
        else:
            stderr = self._noop()
        stdin, stdout, stderr = yield from tasks.gather(stdin, stdout, stderr,
195
                                                        loop=self._loop)
196 197 198 199
        yield from self.wait()
        return (stdout, stderr)


200
@coroutine
201 202 203 204 205 206 207 208 209 210 211 212
def create_subprocess_shell(cmd, stdin=None, stdout=None, stderr=None,
                            loop=None, limit=streams._DEFAULT_LIMIT, **kwds):
    if loop is None:
        loop = events.get_event_loop()
    protocol_factory = lambda: SubprocessStreamProtocol(limit=limit,
                                                        loop=loop)
    transport, protocol = yield from loop.subprocess_shell(
                                            protocol_factory,
                                            cmd, stdin=stdin, stdout=stdout,
                                            stderr=stderr, **kwds)
    return Process(transport, protocol, loop)

213
@coroutine
214 215 216
def create_subprocess_exec(program, *args, stdin=None, stdout=None,
                           stderr=None, loop=None,
                           limit=streams._DEFAULT_LIMIT, **kwds):
217 218 219 220 221 222
    if loop is None:
        loop = events.get_event_loop()
    protocol_factory = lambda: SubprocessStreamProtocol(limit=limit,
                                                        loop=loop)
    transport, protocol = yield from loop.subprocess_exec(
                                            protocol_factory,
223 224
                                            program, *args,
                                            stdin=stdin, stdout=stdout,
225 226
                                            stderr=stderr, **kwds)
    return Process(transport, protocol, loop)