blob: f6d6a141eb8d71397fbe6e5584c9ae7900e0506b [file] [log] [blame]
Victor Stinner915bcb02014-02-01 22:49:59 +01001__all__ = ['create_subprocess_exec', 'create_subprocess_shell']
2
3import collections
4import subprocess
5
6from . import events
7from . import futures
8from . import protocols
9from . import streams
10from . import tasks
Victor Stinnerf951d282014-06-29 00:46:45 +020011from .coroutines import coroutine
Victor Stinneracdb7822014-07-14 18:33:40 +020012from .log import logger
Victor Stinner915bcb02014-02-01 22:49:59 +010013
14
15PIPE = subprocess.PIPE
16STDOUT = subprocess.STDOUT
17DEVNULL = subprocess.DEVNULL
18
19
20class 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
28 self.waiter = futures.Future(loop=loop)
29 self._waiters = collections.deque()
30 self._transport = None
31
Victor Stinneracdb7822014-07-14 18:33:40 +020032 def __repr__(self):
33 info = [self.__class__.__name__]
34 if self.stdin is not None:
35 info.append('stdin=%r' % self.stdin)
36 if self.stdout is not None:
37 info.append('stdout=%r' % self.stdout)
38 if self.stderr is not None:
39 info.append('stderr=%r' % self.stderr)
40 return '<%s>' % ' '.join(info)
41
Victor Stinner915bcb02014-02-01 22:49:59 +010042 def connection_made(self, transport):
43 self._transport = transport
Victor Stinner5ef586f2014-11-25 17:20:33 +010044
45 stdout_transport = transport.get_pipe_transport(1)
46 if stdout_transport is not None:
Victor Stinner915bcb02014-02-01 22:49:59 +010047 self.stdout = streams.StreamReader(limit=self._limit,
48 loop=self._loop)
Victor Stinner5ef586f2014-11-25 17:20:33 +010049 self.stdout.set_transport(stdout_transport)
50
51 stderr_transport = transport.get_pipe_transport(2)
52 if stderr_transport is not None:
Victor Stinner915bcb02014-02-01 22:49:59 +010053 self.stderr = streams.StreamReader(limit=self._limit,
54 loop=self._loop)
Victor Stinner5ef586f2014-11-25 17:20:33 +010055 self.stderr.set_transport(stderr_transport)
56
57 stdin_transport = transport.get_pipe_transport(0)
58 if stdin_transport is not None:
59 self.stdin = streams.StreamWriter(stdin_transport,
Victor Stinner915bcb02014-02-01 22:49:59 +010060 protocol=self,
61 reader=None,
62 loop=self._loop)
63 self.waiter.set_result(None)
64
65 def pipe_data_received(self, fd, data):
66 if fd == 1:
67 reader = self.stdout
68 elif fd == 2:
69 reader = self.stderr
70 else:
71 reader = None
72 if reader is not None:
73 reader.feed_data(data)
74
75 def pipe_connection_lost(self, fd, exc):
76 if fd == 0:
77 pipe = self.stdin
78 if pipe is not None:
79 pipe.close()
80 self.connection_lost(exc)
81 return
82 if fd == 1:
83 reader = self.stdout
84 elif fd == 2:
85 reader = self.stderr
86 else:
87 reader = None
88 if reader != None:
89 if exc is None:
90 reader.feed_eof()
91 else:
92 reader.set_exception(exc)
93
94 def process_exited(self):
95 # wake up futures waiting for wait()
96 returncode = self._transport.get_returncode()
97 while self._waiters:
98 waiter = self._waiters.popleft()
99 waiter.set_result(returncode)
100
101
102class Process:
103 def __init__(self, transport, protocol, loop):
104 self._transport = transport
105 self._protocol = protocol
106 self._loop = loop
107 self.stdin = protocol.stdin
108 self.stdout = protocol.stdout
109 self.stderr = protocol.stderr
110 self.pid = transport.get_pid()
111
Victor Stinneracdb7822014-07-14 18:33:40 +0200112 def __repr__(self):
113 return '<%s %s>' % (self.__class__.__name__, self.pid)
114
Victor Stinner915bcb02014-02-01 22:49:59 +0100115 @property
116 def returncode(self):
117 return self._transport.get_returncode()
118
Victor Stinnerf951d282014-06-29 00:46:45 +0200119 @coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100120 def wait(self):
121 """Wait until the process exit and return the process return code."""
122 returncode = self._transport.get_returncode()
123 if returncode is not None:
124 return returncode
125
126 waiter = futures.Future(loop=self._loop)
127 self._protocol._waiters.append(waiter)
128 yield from waiter
129 return waiter.result()
130
Victor Stinner915bcb02014-02-01 22:49:59 +0100131 def _check_alive(self):
132 if self._transport.get_returncode() is not None:
133 raise ProcessLookupError()
134
135 def send_signal(self, signal):
136 self._check_alive()
137 self._transport.send_signal(signal)
138
139 def terminate(self):
140 self._check_alive()
141 self._transport.terminate()
142
143 def kill(self):
144 self._check_alive()
145 self._transport.kill()
146
Victor Stinnerf951d282014-06-29 00:46:45 +0200147 @coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100148 def _feed_stdin(self, input):
Victor Stinnerd55b54d2014-07-17 13:12:03 +0200149 debug = self._loop.get_debug()
Victor Stinner915bcb02014-02-01 22:49:59 +0100150 self.stdin.write(input)
Victor Stinnerd55b54d2014-07-17 13:12:03 +0200151 if debug:
Victor Stinneracdb7822014-07-14 18:33:40 +0200152 logger.debug('%r communicate: feed stdin (%s bytes)',
153 self, len(input))
Victor Stinnercc996b52014-07-17 12:25:27 +0200154 try:
155 yield from self.stdin.drain()
Victor Stinnerd55b54d2014-07-17 13:12:03 +0200156 except (BrokenPipeError, ConnectionResetError) as exc:
157 # communicate() ignores BrokenPipeError and ConnectionResetError
158 if debug:
159 logger.debug('%r communicate: stdin got %r', self, exc)
Victor Stinneracdb7822014-07-14 18:33:40 +0200160
Victor Stinnerd55b54d2014-07-17 13:12:03 +0200161 if debug:
Victor Stinneracdb7822014-07-14 18:33:40 +0200162 logger.debug('%r communicate: close stdin', self)
Victor Stinner915bcb02014-02-01 22:49:59 +0100163 self.stdin.close()
164
Victor Stinnerf951d282014-06-29 00:46:45 +0200165 @coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100166 def _noop(self):
167 return None
168
Victor Stinnerf951d282014-06-29 00:46:45 +0200169 @coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100170 def _read_stream(self, fd):
171 transport = self._transport.get_pipe_transport(fd)
172 if fd == 2:
173 stream = self.stderr
174 else:
175 assert fd == 1
176 stream = self.stdout
Victor Stinneracdb7822014-07-14 18:33:40 +0200177 if self._loop.get_debug():
178 name = 'stdout' if fd == 1 else 'stderr'
179 logger.debug('%r communicate: read %s', self, name)
Victor Stinner915bcb02014-02-01 22:49:59 +0100180 output = yield from stream.read()
Victor Stinneracdb7822014-07-14 18:33:40 +0200181 if self._loop.get_debug():
182 name = 'stdout' if fd == 1 else 'stderr'
183 logger.debug('%r communicate: close %s', self, name)
Victor Stinner915bcb02014-02-01 22:49:59 +0100184 transport.close()
185 return output
186
Victor Stinnerf951d282014-06-29 00:46:45 +0200187 @coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100188 def communicate(self, input=None):
Victor Stinner915bcb02014-02-01 22:49:59 +0100189 if input:
190 stdin = self._feed_stdin(input)
191 else:
192 stdin = self._noop()
193 if self.stdout is not None:
194 stdout = self._read_stream(1)
195 else:
196 stdout = self._noop()
197 if self.stderr is not None:
198 stderr = self._read_stream(2)
199 else:
200 stderr = self._noop()
201 stdin, stdout, stderr = yield from tasks.gather(stdin, stdout, stderr,
Victor Stinner8d0230b2014-02-20 10:12:59 +0100202 loop=self._loop)
Victor Stinner915bcb02014-02-01 22:49:59 +0100203 yield from self.wait()
204 return (stdout, stderr)
205
206
Victor Stinnerf951d282014-06-29 00:46:45 +0200207@coroutine
Victor Stinner915bcb02014-02-01 22:49:59 +0100208def create_subprocess_shell(cmd, stdin=None, stdout=None, stderr=None,
209 loop=None, limit=streams._DEFAULT_LIMIT, **kwds):
210 if loop is None:
211 loop = events.get_event_loop()
212 protocol_factory = lambda: SubprocessStreamProtocol(limit=limit,
213 loop=loop)
214 transport, protocol = yield from loop.subprocess_shell(
215 protocol_factory,
216 cmd, stdin=stdin, stdout=stdout,
217 stderr=stderr, **kwds)
218 yield from protocol.waiter
219 return Process(transport, protocol, loop)
220
Victor Stinnerf951d282014-06-29 00:46:45 +0200221@coroutine
Yury Selivanov57797522014-02-18 22:56:15 -0500222def create_subprocess_exec(program, *args, stdin=None, stdout=None,
223 stderr=None, loop=None,
224 limit=streams._DEFAULT_LIMIT, **kwds):
Victor Stinner915bcb02014-02-01 22:49:59 +0100225 if loop is None:
226 loop = events.get_event_loop()
227 protocol_factory = lambda: SubprocessStreamProtocol(limit=limit,
228 loop=loop)
229 transport, protocol = yield from loop.subprocess_exec(
230 protocol_factory,
Victor Stinner20e07432014-02-11 11:44:56 +0100231 program, *args,
232 stdin=stdin, stdout=stdout,
Victor Stinner915bcb02014-02-01 22:49:59 +0100233 stderr=stderr, **kwds)
234 yield from protocol.waiter
235 return Process(transport, protocol, loop)