Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 1 | __all__ = ['create_subprocess_exec', 'create_subprocess_shell'] |
| 2 | |
| 3 | import collections |
| 4 | import subprocess |
| 5 | |
| 6 | from . import events |
| 7 | from . import futures |
| 8 | from . import protocols |
| 9 | from . import streams |
| 10 | from . import tasks |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 11 | from .coroutines import coroutine |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 12 | from .log import logger |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 13 | |
| 14 | |
| 15 | PIPE = subprocess.PIPE |
| 16 | STDOUT = subprocess.STDOUT |
| 17 | DEVNULL = subprocess.DEVNULL |
| 18 | |
| 19 | |
| 20 | class SubprocessStreamProtocol(streams.FlowControlMixin, |
| 21 | protocols.SubprocessProtocol): |
| 22 | """Like StreamReaderProtocol, but for a subprocess.""" |
| 23 | |
| 24 | def __init__(self, limit, loop): |
| 25 | super().__init__(loop=loop) |
| 26 | self._limit = limit |
| 27 | self.stdin = self.stdout = self.stderr = None |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 28 | self._transport = None |
| 29 | |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 30 | def __repr__(self): |
| 31 | info = [self.__class__.__name__] |
| 32 | if self.stdin is not None: |
| 33 | info.append('stdin=%r' % self.stdin) |
| 34 | if self.stdout is not None: |
| 35 | info.append('stdout=%r' % self.stdout) |
| 36 | if self.stderr is not None: |
| 37 | info.append('stderr=%r' % self.stderr) |
| 38 | return '<%s>' % ' '.join(info) |
| 39 | |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 40 | def connection_made(self, transport): |
| 41 | self._transport = transport |
Victor Stinner | 5ef586f | 2014-11-25 17:20:33 +0100 | [diff] [blame] | 42 | |
| 43 | stdout_transport = transport.get_pipe_transport(1) |
| 44 | if stdout_transport is not None: |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 45 | self.stdout = streams.StreamReader(limit=self._limit, |
| 46 | loop=self._loop) |
Victor Stinner | 5ef586f | 2014-11-25 17:20:33 +0100 | [diff] [blame] | 47 | self.stdout.set_transport(stdout_transport) |
| 48 | |
| 49 | stderr_transport = transport.get_pipe_transport(2) |
| 50 | if stderr_transport is not None: |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 51 | self.stderr = streams.StreamReader(limit=self._limit, |
| 52 | loop=self._loop) |
Victor Stinner | 5ef586f | 2014-11-25 17:20:33 +0100 | [diff] [blame] | 53 | self.stderr.set_transport(stderr_transport) |
| 54 | |
| 55 | stdin_transport = transport.get_pipe_transport(0) |
| 56 | if stdin_transport is not None: |
| 57 | self.stdin = streams.StreamWriter(stdin_transport, |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 58 | protocol=self, |
| 59 | reader=None, |
| 60 | loop=self._loop) |
Victor Stinner | f651a60 | 2015-01-14 02:10:33 +0100 | [diff] [blame] | 61 | |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 62 | def pipe_data_received(self, fd, data): |
| 63 | if fd == 1: |
| 64 | reader = self.stdout |
| 65 | elif fd == 2: |
| 66 | reader = self.stderr |
| 67 | else: |
| 68 | reader = None |
| 69 | if reader is not None: |
| 70 | reader.feed_data(data) |
| 71 | |
| 72 | def pipe_connection_lost(self, fd, exc): |
| 73 | if fd == 0: |
| 74 | pipe = self.stdin |
| 75 | if pipe is not None: |
| 76 | pipe.close() |
| 77 | self.connection_lost(exc) |
| 78 | return |
| 79 | if fd == 1: |
| 80 | reader = self.stdout |
| 81 | elif fd == 2: |
| 82 | reader = self.stderr |
| 83 | else: |
| 84 | reader = None |
| 85 | if reader != None: |
| 86 | if exc is None: |
| 87 | reader.feed_eof() |
| 88 | else: |
| 89 | reader.set_exception(exc) |
| 90 | |
| 91 | def process_exited(self): |
Victor Stinner | 791009b | 2015-01-15 13:16:02 +0100 | [diff] [blame] | 92 | self._transport.close() |
| 93 | self._transport = None |
| 94 | |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 95 | |
| 96 | class Process: |
| 97 | def __init__(self, transport, protocol, loop): |
| 98 | self._transport = transport |
| 99 | self._protocol = protocol |
| 100 | self._loop = loop |
| 101 | self.stdin = protocol.stdin |
| 102 | self.stdout = protocol.stdout |
| 103 | self.stderr = protocol.stderr |
| 104 | self.pid = transport.get_pid() |
| 105 | |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 106 | def __repr__(self): |
| 107 | return '<%s %s>' % (self.__class__.__name__, self.pid) |
| 108 | |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 109 | @property |
| 110 | def returncode(self): |
| 111 | return self._transport.get_returncode() |
| 112 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 113 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 114 | def wait(self): |
Victor Stinner | 47cd10d | 2015-01-30 00:05:19 +0100 | [diff] [blame] | 115 | """Wait until the process exit and return the process return code. |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 116 | |
Victor Stinner | 47cd10d | 2015-01-30 00:05:19 +0100 | [diff] [blame] | 117 | This method is a coroutine.""" |
| 118 | return (yield from self._transport.wait()) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 119 | |
| 120 | def send_signal(self, signal): |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 121 | self._transport.send_signal(signal) |
| 122 | |
| 123 | def terminate(self): |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 124 | self._transport.terminate() |
| 125 | |
| 126 | def kill(self): |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 127 | self._transport.kill() |
| 128 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 129 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 130 | def _feed_stdin(self, input): |
Victor Stinner | d55b54d | 2014-07-17 13:12:03 +0200 | [diff] [blame] | 131 | debug = self._loop.get_debug() |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 132 | self.stdin.write(input) |
Victor Stinner | d55b54d | 2014-07-17 13:12:03 +0200 | [diff] [blame] | 133 | if debug: |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 134 | logger.debug('%r communicate: feed stdin (%s bytes)', |
| 135 | self, len(input)) |
Victor Stinner | cc996b5 | 2014-07-17 12:25:27 +0200 | [diff] [blame] | 136 | try: |
| 137 | yield from self.stdin.drain() |
Victor Stinner | d55b54d | 2014-07-17 13:12:03 +0200 | [diff] [blame] | 138 | except (BrokenPipeError, ConnectionResetError) as exc: |
| 139 | # communicate() ignores BrokenPipeError and ConnectionResetError |
| 140 | if debug: |
| 141 | logger.debug('%r communicate: stdin got %r', self, exc) |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 142 | |
Victor Stinner | d55b54d | 2014-07-17 13:12:03 +0200 | [diff] [blame] | 143 | if debug: |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 144 | logger.debug('%r communicate: close stdin', self) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 145 | self.stdin.close() |
| 146 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 147 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 148 | def _noop(self): |
| 149 | return None |
| 150 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 151 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 152 | def _read_stream(self, fd): |
| 153 | transport = self._transport.get_pipe_transport(fd) |
| 154 | if fd == 2: |
| 155 | stream = self.stderr |
| 156 | else: |
| 157 | assert fd == 1 |
| 158 | stream = self.stdout |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 159 | if self._loop.get_debug(): |
| 160 | name = 'stdout' if fd == 1 else 'stderr' |
| 161 | logger.debug('%r communicate: read %s', self, name) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 162 | output = yield from stream.read() |
Victor Stinner | acdb782 | 2014-07-14 18:33:40 +0200 | [diff] [blame] | 163 | if self._loop.get_debug(): |
| 164 | name = 'stdout' if fd == 1 else 'stderr' |
| 165 | logger.debug('%r communicate: close %s', self, name) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 166 | transport.close() |
| 167 | return output |
| 168 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 169 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 170 | def communicate(self, input=None): |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 171 | if input: |
| 172 | stdin = self._feed_stdin(input) |
| 173 | else: |
| 174 | stdin = self._noop() |
| 175 | if self.stdout is not None: |
| 176 | stdout = self._read_stream(1) |
| 177 | else: |
| 178 | stdout = self._noop() |
| 179 | if self.stderr is not None: |
| 180 | stderr = self._read_stream(2) |
| 181 | else: |
| 182 | stderr = self._noop() |
| 183 | stdin, stdout, stderr = yield from tasks.gather(stdin, stdout, stderr, |
Victor Stinner | 8d0230b | 2014-02-20 10:12:59 +0100 | [diff] [blame] | 184 | loop=self._loop) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 185 | yield from self.wait() |
| 186 | return (stdout, stderr) |
| 187 | |
| 188 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 189 | @coroutine |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 190 | def create_subprocess_shell(cmd, stdin=None, stdout=None, stderr=None, |
| 191 | loop=None, limit=streams._DEFAULT_LIMIT, **kwds): |
| 192 | if loop is None: |
| 193 | loop = events.get_event_loop() |
| 194 | protocol_factory = lambda: SubprocessStreamProtocol(limit=limit, |
| 195 | loop=loop) |
| 196 | transport, protocol = yield from loop.subprocess_shell( |
| 197 | protocol_factory, |
| 198 | cmd, stdin=stdin, stdout=stdout, |
| 199 | stderr=stderr, **kwds) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 200 | return Process(transport, protocol, loop) |
| 201 | |
Victor Stinner | f951d28 | 2014-06-29 00:46:45 +0200 | [diff] [blame] | 202 | @coroutine |
Yury Selivanov | 5779752 | 2014-02-18 22:56:15 -0500 | [diff] [blame] | 203 | def create_subprocess_exec(program, *args, stdin=None, stdout=None, |
| 204 | stderr=None, loop=None, |
| 205 | limit=streams._DEFAULT_LIMIT, **kwds): |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 206 | if loop is None: |
| 207 | loop = events.get_event_loop() |
| 208 | protocol_factory = lambda: SubprocessStreamProtocol(limit=limit, |
| 209 | loop=loop) |
| 210 | transport, protocol = yield from loop.subprocess_exec( |
| 211 | protocol_factory, |
Victor Stinner | 20e0743 | 2014-02-11 11:44:56 +0100 | [diff] [blame] | 212 | program, *args, |
| 213 | stdin=stdin, stdout=stdout, |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 214 | stderr=stderr, **kwds) |
Victor Stinner | 915bcb0 | 2014-02-01 22:49:59 +0100 | [diff] [blame] | 215 | return Process(transport, protocol, loop) |