Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Magic functions for running cells in various scripts."""
3# Copyright (c) IPython Development Team.
4# Distributed under the terms of the Modified BSD License.
6import asyncio
7import asyncio.exceptions
8import atexit
9import errno
10import os
11import signal
12import sys
13import time
14import weakref
15from codecs import getincrementaldecoder
16from subprocess import CalledProcessError
17from threading import Thread
19from traitlets import Any, Dict, List, default
21from IPython.core import magic_arguments
22from IPython.core.async_helpers import _AsyncIOProxy
23from IPython.core.magic import Magics, cell_magic, line_magic, magics_class
24from IPython.utils.process import arg_split
26#-----------------------------------------------------------------------------
27# Magic implementation classes
28#-----------------------------------------------------------------------------
30def script_args(f):
31 """single decorator for adding script args"""
32 args = [
33 magic_arguments.argument(
34 '--out', type=str,
35 help="""The variable in which to store stdout from the script.
36 If the script is backgrounded, this will be the stdout *pipe*,
37 instead of the stderr text itself and will not be auto closed.
38 """
39 ),
40 magic_arguments.argument(
41 '--err', type=str,
42 help="""The variable in which to store stderr from the script.
43 If the script is backgrounded, this will be the stderr *pipe*,
44 instead of the stderr text itself and will not be autoclosed.
45 """
46 ),
47 magic_arguments.argument(
48 '--bg', action="store_true",
49 help="""Whether to run the script in the background.
50 If given, the only way to see the output of the command is
51 with --out/err.
52 """
53 ),
54 magic_arguments.argument(
55 '--proc', type=str,
56 help="""The variable in which to store Popen instance.
57 This is used only when --bg option is given.
58 """
59 ),
60 magic_arguments.argument(
61 '--no-raise-error', action="store_false", dest='raise_error',
62 help="""Whether you should raise an error message in addition to
63 a stream on stderr if you get a nonzero exit code.
64 """,
65 ),
66 ]
67 for arg in args:
68 f = arg(f)
69 return f
72class RaiseAfterInterrupt(Exception):
73 pass
76@magics_class
77class ScriptMagics(Magics):
78 """Magics for talking to scripts
80 This defines a base `%%script` cell magic for running a cell
81 with a program in a subprocess, and registers a few top-level
82 magics that call %%script with common interpreters.
83 """
85 event_loop = Any(
86 help="""
87 The event loop on which to run subprocesses
89 Not the main event loop,
90 because we want to be able to make blocking calls
91 and have certain requirements we don't want to impose on the main loop.
92 """
93 )
95 script_magics: List = List(
96 help="""Extra script cell magics to define
98 This generates simple wrappers of `%%script foo` as `%%foo`.
100 If you want to add script magics that aren't on your path,
101 specify them in script_paths
102 """,
103 ).tag(config=True)
105 @default('script_magics')
106 def _script_magics_default(self):
107 """default to a common list of programs"""
109 defaults = [
110 'sh',
111 'bash',
112 'perl',
113 'ruby',
114 'python',
115 'python2',
116 'python3',
117 'pypy',
118 ]
119 if os.name == 'nt':
120 defaults.extend([
121 'cmd',
122 ])
124 return defaults
126 script_paths = Dict(
127 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
129 Only necessary for items in script_magics where the default path will not
130 find the right interpreter.
131 """
132 ).tag(config=True)
134 def __init__(self, shell=None):
135 super().__init__(shell=shell)
136 self._generate_script_magics()
137 self.bg_processes = []
138 self._event_loop_finalizer = None
139 atexit.register(self.kill_bg_processes)
141 def __del__(self):
142 self.kill_bg_processes()
144 @staticmethod
145 def _shutdown_event_loop(event_loop, thread):
146 """Stop ``event_loop``, wait for ``thread`` to notice, and close it.
148 Kept free of any reference to the ``ScriptMagics`` instance so it can
149 be handed to :func:`weakref.finalize` without keeping that instance
150 alive.
151 """
152 if not event_loop.is_closed():
153 event_loop.call_soon_threadsafe(event_loop.stop)
154 thread.join()
155 event_loop.close()
157 def stop_event_loop(self):
158 """Stop the background event loop and the thread running it.
160 The loop is started lazily by ``shebang`` and then kept around to be
161 reused; this is the deterministic way to shut it back down. Without it
162 the thread lives until the process exits, which leaves the loop (and
163 the socketpair it uses for its self-pipe) unclosed, and keeps the
164 process multi-threaded, which ``os.fork()`` warns about since
165 Python 3.12. Safe to call more than once.
167 The same shutdown runs on its own if this object is garbage collected,
168 and at interpreter exit, through the finalizer ``shebang`` registers.
169 """
170 finalizer, self._event_loop_finalizer = self._event_loop_finalizer, None
171 self.event_loop = None
172 if finalizer is not None:
173 finalizer()
175 def _generate_script_magics(self):
176 cell_magics = self.magics['cell']
177 for name in self.script_magics:
178 cell_magics[name] = self._make_script_magic(name)
180 def _make_script_magic(self, name):
181 """make a named magic, that calls %%script with a particular program"""
182 # expand to explicit path if necessary:
183 script = self.script_paths.get(name, name)
185 @magic_arguments.magic_arguments()
186 @script_args
187 def named_script_magic(line, cell):
188 # if line, add it as cl-flags
189 if line:
190 line = "{} {}".format(script, line)
191 else:
192 line = script
193 return self.shebang(line, cell)
195 # write a basic docstring:
196 named_script_magic.__doc__ = f"""%%{name} script magic
198 Run cells with {script} in a subprocess.
200 This is a shortcut for `%%script {script}`
201 """
203 return named_script_magic
205 @magic_arguments.magic_arguments()
206 @script_args
207 @cell_magic("script")
208 def shebang(self, line, cell):
209 """Run a cell via a shell command
211 The `%%script` line is like the #! line of script,
212 specifying a program (bash, perl, ruby, etc.) with which to run.
214 The rest of the cell is run by that program.
216 .. versionchanged:: 9.0
217 Interrupting the script executed without `--bg` will end in
218 raising an exception (unless `--no-raise-error` is passed).
220 Examples
221 --------
222 ::
224 In [1]: %%script bash
225 ...: for i in 1 2 3; do
226 ...: echo $i
227 ...: done
228 1
229 2
230 3
231 """
233 # Create the event loop in which to run script magics
234 # this operates on a background thread
235 if self.event_loop is None:
236 if sys.platform == "win32":
237 # don't override the current policy,
238 # just create an event loop
239 event_loop = asyncio.WindowsProactorEventLoopPolicy().new_event_loop()
240 else:
241 event_loop = asyncio.new_event_loop()
242 self.event_loop = event_loop
244 # start the loop in a background thread
245 asyncio_thread = Thread(target=event_loop.run_forever, daemon=True)
246 asyncio_thread.start()
247 # ... and make sure it is stopped again, at the latest when we are
248 # collected or the interpreter exits
249 self._event_loop_finalizer = weakref.finalize(
250 self, self._shutdown_event_loop, event_loop, asyncio_thread
251 )
252 else:
253 event_loop = self.event_loop
255 def in_thread(coro):
256 """Call a coroutine on the asyncio thread"""
257 return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
259 async def _readchunk(stream):
260 try:
261 return await stream.read(100)
262 except asyncio.exceptions.IncompleteReadError as e:
263 return e.partial
264 except asyncio.exceptions.LimitOverrunError as e:
265 return await stream.read(e.consumed)
267 async def _handle_stream(stream, stream_arg, file_object):
268 should_break = False
269 decoder = getincrementaldecoder("utf-8")(errors="replace")
270 while True:
271 chunk = decoder.decode(await _readchunk(stream))
272 if not chunk:
273 break
274 chunk = decoder.decode("", final=True)
275 should_break = True
276 if stream_arg:
277 self.shell.user_ns[stream_arg] += chunk
278 else:
279 file_object.write(chunk)
280 file_object.flush()
281 if should_break:
282 break
284 async def _stream_communicate(process, cell):
285 process.stdin.write(cell)
286 process.stdin.close()
287 stdout_task = asyncio.create_task(
288 _handle_stream(process.stdout, args.out, sys.stdout)
289 )
290 stderr_task = asyncio.create_task(
291 _handle_stream(process.stderr, args.err, sys.stderr)
292 )
293 await asyncio.wait([stdout_task, stderr_task])
294 await process.wait()
296 argv = arg_split(line, posix=not sys.platform.startswith("win"))
297 args, cmd = self.shebang.parser.parse_known_args(argv)
299 if args.out:
300 self.shell.user_ns[args.out] = ""
301 if args.err:
302 self.shell.user_ns[args.err] = ""
304 try:
305 p = in_thread(
306 asyncio.create_subprocess_exec(
307 *cmd,
308 stdout=asyncio.subprocess.PIPE,
309 stderr=asyncio.subprocess.PIPE,
310 stdin=asyncio.subprocess.PIPE,
311 )
312 )
313 except OSError as e:
314 if e.errno == errno.ENOENT:
315 print("Couldn't find program: %r" % cmd[0])
316 return
317 else:
318 raise
320 if not cell.endswith('\n'):
321 cell += '\n'
322 cell = cell.encode('utf8', 'replace')
323 if args.bg:
324 self.bg_processes.append(p)
325 self._gc_bg_processes()
326 to_close = []
327 if args.out:
328 self.shell.user_ns[args.out] = _AsyncIOProxy(p.stdout, event_loop)
329 else:
330 to_close.append(p.stdout)
331 if args.err:
332 self.shell.user_ns[args.err] = _AsyncIOProxy(p.stderr, event_loop)
333 else:
334 to_close.append(p.stderr)
335 event_loop.call_soon_threadsafe(
336 lambda: asyncio.Task(self._run_script(p, cell, to_close))
337 )
338 if args.proc:
339 proc_proxy = _AsyncIOProxy(p, event_loop)
340 proc_proxy.stdout = _AsyncIOProxy(p.stdout, event_loop)
341 proc_proxy.stderr = _AsyncIOProxy(p.stderr, event_loop)
342 self.shell.user_ns[args.proc] = proc_proxy
343 return
345 try:
346 in_thread(_stream_communicate(p, cell))
347 except KeyboardInterrupt:
348 try:
349 p.send_signal(signal.SIGINT)
350 in_thread(asyncio.wait_for(p.wait(), timeout=0.1))
351 if p.returncode is not None:
352 print("Process was interrupted.")
353 if args.raise_error:
354 raise RaiseAfterInterrupt()
355 else:
356 return
357 p.terminate()
358 in_thread(asyncio.wait_for(p.wait(), timeout=0.1))
359 if p.returncode is not None:
360 print("Process was terminated.")
361 if args.raise_error:
362 raise RaiseAfterInterrupt()
363 else:
364 return
365 p.kill()
366 print("Process was killed.")
367 if args.raise_error:
368 raise RaiseAfterInterrupt()
369 except RaiseAfterInterrupt:
370 pass
371 except OSError:
372 pass
373 except Exception as e:
374 print("Error while terminating subprocess (pid=%i): %s" % (p.pid, e))
375 if args.raise_error:
376 raise CalledProcessError(p.returncode, cell) from None
377 else:
378 return
380 if args.raise_error and p.returncode != 0:
381 # If we get here and p.returncode is still None, we must have
382 # killed it but not yet seen its return code. We don't wait for it,
383 # in case it's stuck in uninterruptible sleep. -9 = SIGKILL
384 rc = p.returncode or -9
385 raise CalledProcessError(rc, cell)
387 shebang.__skip_doctest__ = os.name != "posix"
389 async def _run_script(self, p, cell, to_close):
390 """callback for running the script in the background"""
392 p.stdin.write(cell)
393 await p.stdin.drain()
394 p.stdin.close()
395 await p.stdin.wait_closed()
396 await p.wait()
397 # asyncio read pipes have no close
398 # but we should drain the data anyway
399 for s in to_close:
400 await s.read()
401 self._gc_bg_processes()
403 @line_magic("killbgscripts")
404 def killbgscripts(self, _nouse_=''):
405 """Kill all BG processes started by %%script and its family."""
406 self.kill_bg_processes()
407 print("All background processes were killed.")
409 def kill_bg_processes(self):
410 """Kill all BG processes which are still running."""
411 if not self.bg_processes:
412 return
413 for p in self.bg_processes:
414 if p.returncode is None:
415 try:
416 p.send_signal(signal.SIGINT)
417 except OSError:
418 pass
419 time.sleep(0.1)
420 self._gc_bg_processes()
421 if not self.bg_processes:
422 return
423 for p in self.bg_processes:
424 if p.returncode is None:
425 try:
426 p.terminate()
427 except OSError:
428 pass
429 time.sleep(0.1)
430 self._gc_bg_processes()
431 if not self.bg_processes:
432 return
433 for p in self.bg_processes:
434 if p.returncode is None:
435 try:
436 p.kill()
437 except OSError:
438 pass
439 self._gc_bg_processes()
441 def _gc_bg_processes(self):
442 self.bg_processes = [p for p in self.bg_processes if p.returncode is None]