##// END OF EJS Templates
Line-buffered output for %%magic scripts
Ethan Madden -
Show More
@@ -8,7 +8,7 b' import os'
8 8 import sys
9 9 import signal
10 10 import time
11 from subprocess import Popen, PIPE, CalledProcessError
11 import asyncio
12 12 import atexit
13 13
14 14 from IPython.core import magic_arguments
@@ -175,11 +175,44 b' class ScriptMagics(Magics):'
175 175 2
176 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 205 argv = arg_split(line, posix = not sys.platform.startswith('win'))
179 206 args, cmd = self.shebang.parser.parse_known_args(argv)
180
181 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 216 except OSError as e:
184 217 if e.errno == errno.ENOENT:
185 218 print("Couldn't find program: %r" % cmd[0])
@@ -208,17 +241,17 b' class ScriptMagics(Magics):'
208 241 return
209 242
210 243 try:
211 out, err = p.communicate(cell)
244 loop.run_until_complete(_stream_communicate(p, cell))
212 245 except KeyboardInterrupt:
213 246 try:
214 247 p.send_signal(signal.SIGINT)
215 248 time.sleep(0.1)
216 if p.poll() is not None:
249 if p.returncode is not None:
217 250 print("Process is interrupted.")
218 251 return
219 252 p.terminate()
220 253 time.sleep(0.1)
221 if p.poll() is not None:
254 if p.returncode is not None:
222 255 print("Process is terminated.")
223 256 return
224 257 p.kill()
@@ -229,20 +262,9 b' class ScriptMagics(Magics):'
229 262 print("Error while terminating subprocess (pid=%i): %s" \
230 263 % (p.pid, e))
231 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 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 269 def _run_script(self, p, cell, to_close):
248 270 """callback for running the script in the background"""
@@ -263,7 +285,7 b' class ScriptMagics(Magics):'
263 285 if not self.bg_processes:
264 286 return
265 287 for p in self.bg_processes:
266 if p.poll() is None:
288 if p.returncode is None:
267 289 try:
268 290 p.send_signal(signal.SIGINT)
269 291 except:
@@ -273,7 +295,7 b' class ScriptMagics(Magics):'
273 295 if not self.bg_processes:
274 296 return
275 297 for p in self.bg_processes:
276 if p.poll() is None:
298 if p.returncode is None:
277 299 try:
278 300 p.terminate()
279 301 except:
@@ -283,7 +305,7 b' class ScriptMagics(Magics):'
283 305 if not self.bg_processes:
284 306 return
285 307 for p in self.bg_processes:
286 if p.poll() is None:
308 if p.returncode is None:
287 309 try:
288 310 p.kill()
289 311 except:
@@ -291,4 +313,4 b' class ScriptMagics(Magics):'
291 313 self._gc_bg_processes()
292 314
293 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 948 nt.assert_equal(ip.user_ns['error'], 'hello\n')
949 949
950 950 @dec.skip_win32
951 def test_script_bg_out():
951 async def test_script_bg_out():
952 952 ip = get_ipython()
953 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 956 ip.user_ns['output'].close()
957 957
958 958 @dec.skip_win32
959 def test_script_bg_err():
959 async def test_script_bg_err():
960 960 ip = get_ipython()
961 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 963 ip.user_ns['error'].close()
964 964
965 965 @dec.skip_win32
966 def test_script_bg_out_err():
966 async def test_script_bg_out_err():
967 967 ip = get_ipython()
968 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')
970 nt.assert_equal(ip.user_ns['error'].read(), b'hello\n')
969 nt.assert_equal((await ip.user_ns['output'].read()), b'hi\n')
970 nt.assert_equal((await ip.user_ns['error'].read()), b'hello\n')
971 971 ip.user_ns['output'].close()
972 972 ip.user_ns['error'].close()
973 973
General Comments 0
You need to be logged in to leave comments. Login now