##// END OF EJS Templates
Line-buffered output for %%magic scripts
Ethan Madden -
Show More
@@ -8,7 +8,7 b' import os'
8 import sys
8 import sys
9 import signal
9 import signal
10 import time
10 import time
11 from subprocess import Popen, PIPE, CalledProcessError
11 import asyncio
12 import atexit
12 import atexit
13
13
14 from IPython.core import magic_arguments
14 from IPython.core import magic_arguments
@@ -175,11 +175,44 b' class ScriptMagics(Magics):'
175 2
175 2
176 3
176 3
177 """
177 """
178
179 async def _handle_stream(stream, stream_arg, file_object):
180 while True:
181 line = (await stream.readline()).decode("utf8")
182 if not line:
183 break
184 if stream_arg:
185 self.shell.user_ns[stream_arg] = line
186 else:
187 file_object.write(line)
188 file_object.flush()
189
190 async def _stream_communicate(process, cell):
191 process.stdin.write(cell)
192 process.stdin.close()
193 stdout_task = asyncio.create_task(
194 _handle_stream(process.stdout, args.out, sys.stdout)
195 )
196 stderr_task = asyncio.create_task(
197 _handle_stream(process.stderr, args.err, sys.stderr)
198 )
199 await asyncio.wait([stdout_task, stderr_task])
200
201 if sys.platform.startswith("win"):
202 asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
203 loop = asyncio.get_event_loop()
204
178 argv = arg_split(line, posix = not sys.platform.startswith('win'))
205 argv = arg_split(line, posix=not sys.platform.startswith('win'))
179 args, cmd = self.shebang.parser.parse_known_args(argv)
206 args, cmd = self.shebang.parser.parse_known_args(argv)
180
181 try:
207 try:
182 p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE)
208 p = loop.run_until_complete(
209 asyncio.create_subprocess_exec(
210 *cmd,
211 stdout=asyncio.subprocess.PIPE,
212 stderr=asyncio.subprocess.PIPE,
213 stdin=asyncio.subprocess.PIPE
214 )
215 )
183 except OSError as e:
216 except OSError as e:
184 if e.errno == errno.ENOENT:
217 if e.errno == errno.ENOENT:
185 print("Couldn't find program: %r" % cmd[0])
218 print("Couldn't find program: %r" % cmd[0])
@@ -208,17 +241,17 b' class ScriptMagics(Magics):'
208 return
241 return
209
242
210 try:
243 try:
211 out, err = p.communicate(cell)
244 loop.run_until_complete(_stream_communicate(p, cell))
212 except KeyboardInterrupt:
245 except KeyboardInterrupt:
213 try:
246 try:
214 p.send_signal(signal.SIGINT)
247 p.send_signal(signal.SIGINT)
215 time.sleep(0.1)
248 time.sleep(0.1)
216 if p.poll() is not None:
249 if p.returncode is not None:
217 print("Process is interrupted.")
250 print("Process is interrupted.")
218 return
251 return
219 p.terminate()
252 p.terminate()
220 time.sleep(0.1)
253 time.sleep(0.1)
221 if p.poll() is not None:
254 if p.returncode is not None:
222 print("Process is terminated.")
255 print("Process is terminated.")
223 return
256 return
224 p.kill()
257 p.kill()
@@ -229,20 +262,9 b' class ScriptMagics(Magics):'
229 print("Error while terminating subprocess (pid=%i): %s" \
262 print("Error while terminating subprocess (pid=%i): %s" \
230 % (p.pid, e))
263 % (p.pid, e))
231 return
264 return
232 out = py3compat.decode(out)
233 err = py3compat.decode(err)
234 if args.out:
235 self.shell.user_ns[args.out] = out
236 else:
237 sys.stdout.write(out)
238 sys.stdout.flush()
239 if args.err:
240 self.shell.user_ns[args.err] = err
241 else:
242 sys.stderr.write(err)
243 sys.stderr.flush()
244 if args.raise_error and p.returncode!=0:
265 if args.raise_error and p.returncode!=0:
245 raise CalledProcessError(p.returncode, cell, output=out, stderr=err)
266 print(p.returncode)
267 raise CalledProcessError(p.returncode, cell)
246
268
247 def _run_script(self, p, cell, to_close):
269 def _run_script(self, p, cell, to_close):
248 """callback for running the script in the background"""
270 """callback for running the script in the background"""
@@ -263,7 +285,7 b' class ScriptMagics(Magics):'
263 if not self.bg_processes:
285 if not self.bg_processes:
264 return
286 return
265 for p in self.bg_processes:
287 for p in self.bg_processes:
266 if p.poll() is None:
288 if p.returncode is None:
267 try:
289 try:
268 p.send_signal(signal.SIGINT)
290 p.send_signal(signal.SIGINT)
269 except:
291 except:
@@ -273,7 +295,7 b' class ScriptMagics(Magics):'
273 if not self.bg_processes:
295 if not self.bg_processes:
274 return
296 return
275 for p in self.bg_processes:
297 for p in self.bg_processes:
276 if p.poll() is None:
298 if p.returncode is None:
277 try:
299 try:
278 p.terminate()
300 p.terminate()
279 except:
301 except:
@@ -283,7 +305,7 b' class ScriptMagics(Magics):'
283 if not self.bg_processes:
305 if not self.bg_processes:
284 return
306 return
285 for p in self.bg_processes:
307 for p in self.bg_processes:
286 if p.poll() is None:
308 if p.returncode is None:
287 try:
309 try:
288 p.kill()
310 p.kill()
289 except:
311 except:
@@ -291,4 +313,4 b' class ScriptMagics(Magics):'
291 self._gc_bg_processes()
313 self._gc_bg_processes()
292
314
293 def _gc_bg_processes(self):
315 def _gc_bg_processes(self):
294 self.bg_processes = [p for p in self.bg_processes if p.poll() is None]
316 self.bg_processes = [p for p in self.bg_processes if p.returncode is None]
@@ -948,26 +948,26 b' def test_script_out_err():'
948 nt.assert_equal(ip.user_ns['error'], 'hello\n')
948 nt.assert_equal(ip.user_ns['error'], 'hello\n')
949
949
950 @dec.skip_win32
950 @dec.skip_win32
951 def test_script_bg_out():
951 async def test_script_bg_out():
952 ip = get_ipython()
952 ip = get_ipython()
953 ip.run_cell_magic("script", "--bg --out output sh", "echo 'hi'")
953 ip.run_cell_magic("script", "--bg --out output sh", "echo 'hi'")
954
954
955 nt.assert_equal(ip.user_ns['output'].read(), b'hi\n')
955 nt.assert_equal((await ip.user_ns['output'].read()), b'hi\n')
956 ip.user_ns['output'].close()
956 ip.user_ns['output'].close()
957
957
958 @dec.skip_win32
958 @dec.skip_win32
959 def test_script_bg_err():
959 async def test_script_bg_err():
960 ip = get_ipython()
960 ip = get_ipython()
961 ip.run_cell_magic("script", "--bg --err error sh", "echo 'hello' >&2")
961 ip.run_cell_magic("script", "--bg --err error sh", "echo 'hello' >&2")
962 nt.assert_equal(ip.user_ns['error'].read(), b'hello\n')
962 nt.assert_equal((await ip.user_ns['error'].read()), b'hello\n')
963 ip.user_ns['error'].close()
963 ip.user_ns['error'].close()
964
964
965 @dec.skip_win32
965 @dec.skip_win32
966 def test_script_bg_out_err():
966 async def test_script_bg_out_err():
967 ip = get_ipython()
967 ip = get_ipython()
968 ip.run_cell_magic("script", "--bg --out output --err error sh", "echo 'hi'\necho 'hello' >&2")
968 ip.run_cell_magic("script", "--bg --out output --err error sh", "echo 'hi'\necho 'hello' >&2")
969 nt.assert_equal(ip.user_ns['output'].read(), b'hi\n')
969 nt.assert_equal((await ip.user_ns['output'].read()), b'hi\n')
970 nt.assert_equal(ip.user_ns['error'].read(), b'hello\n')
970 nt.assert_equal((await ip.user_ns['error'].read()), b'hello\n')
971 ip.user_ns['output'].close()
971 ip.user_ns['output'].close()
972 ip.user_ns['error'].close()
972 ip.user_ns['error'].close()
973
973
General Comments 0
You need to be logged in to leave comments. Login now