script.py
316 lines
| 10.1 KiB
| text/x-python
|
PythonLexer
|
r7299 | """Magic functions for running cells in various scripts.""" | ||
|
r22340 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
|
r7299 | |||
|
r8470 | import errno | ||
|
r7299 | import os | ||
import sys | ||||
|
r7561 | import signal | ||
import time | ||||
|
r26199 | import asyncio | ||
|
r7622 | import atexit | ||
|
r7299 | |||
|
r26207 | from subprocess import CalledProcessError | ||
|
r7398 | from IPython.core import magic_arguments | ||
|
r7299 | from IPython.core.magic import ( | ||
Magics, magics_class, line_magic, cell_magic | ||||
) | ||||
|
r7398 | from IPython.lib.backgroundjobs import BackgroundJobManager | ||
|
r8469 | from IPython.utils.process import arg_split | ||
|
r22340 | from traitlets import List, Dict, default | ||
|
r7299 | |||
|
r26207 | |||
|
r7299 | #----------------------------------------------------------------------------- | ||
# Magic implementation classes | ||||
#----------------------------------------------------------------------------- | ||||
|
r7398 | def script_args(f): | ||
"""single decorator for adding script args""" | ||||
args = [ | ||||
magic_arguments.argument( | ||||
'--out', type=str, | ||||
help="""The variable in which to store stdout from the script. | ||||
If the script is backgrounded, this will be the stdout *pipe*, | ||||
|
r24375 | instead of the stderr text itself and will not be auto closed. | ||
|
r7398 | """ | ||
), | ||||
magic_arguments.argument( | ||||
'--err', type=str, | ||||
help="""The variable in which to store stderr from the script. | ||||
If the script is backgrounded, this will be the stderr *pipe*, | ||||
|
r24375 | instead of the stderr text itself and will not be autoclosed. | ||
|
r7398 | """ | ||
), | ||||
magic_arguments.argument( | ||||
'--bg', action="store_true", | ||||
help="""Whether to run the script in the background. | ||||
If given, the only way to see the output of the command is | ||||
with --out/err. | ||||
""" | ||||
), | ||||
|
r7549 | magic_arguments.argument( | ||
'--proc', type=str, | ||||
help="""The variable in which to store Popen instance. | ||||
This is used only when --bg option is given. | ||||
""" | ||||
), | ||||
|
r24410 | magic_arguments.argument( | ||
|
r24519 | '--no-raise-error', action="store_false", dest='raise_error', | ||
|
r24410 | help="""Whether you should raise an error message in addition to | ||
a stream on stderr if you get a nonzero exit code. | ||||
""" | ||||
) | ||||
|
r7398 | ] | ||
for arg in args: | ||||
f = arg(f) | ||||
return f | ||||
|
r7299 | @magics_class | ||
|
r13237 | class ScriptMagics(Magics): | ||
|
r7299 | """Magics for talking to scripts | ||
This defines a base `%%script` cell magic for running a cell | ||||
with a program in a subprocess, and registers a few top-level | ||||
magics that call %%script with common interpreters. | ||||
""" | ||||
|
r22340 | script_magics = List( | ||
|
r7299 | help="""Extra script cell magics to define | ||
This generates simple wrappers of `%%script foo` as `%%foo`. | ||||
If you want to add script magics that aren't on your path, | ||||
specify them in script_paths | ||||
""", | ||||
|
r22340 | ).tag(config=True) | ||
@default('script_magics') | ||||
|
r7299 | def _script_magics_default(self): | ||
|
r8469 | """default to a common list of programs""" | ||
|
r7299 | |||
|
r8469 | defaults = [ | ||
|
r7299 | 'sh', | ||
'bash', | ||||
'perl', | ||||
'ruby', | ||||
|
r8469 | 'python', | ||
|
r16020 | 'python2', | ||
|
r7299 | 'python3', | ||
'pypy', | ||||
|
r8469 | ] | ||
if os.name == 'nt': | ||||
defaults.extend([ | ||||
'cmd', | ||||
]) | ||||
|
r7299 | |||
return defaults | ||||
|
r22340 | script_paths = Dict( | ||
|
r7299 | help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby' | ||
Only necessary for items in script_magics where the default path will not | ||||
find the right interpreter. | ||||
""" | ||||
|
r22340 | ).tag(config=True) | ||
|
r7299 | |||
def __init__(self, shell=None): | ||||
|
r13237 | super(ScriptMagics, self).__init__(shell=shell) | ||
|
r7299 | self._generate_script_magics() | ||
|
r7398 | self.job_manager = BackgroundJobManager() | ||
|
r7618 | self.bg_processes = [] | ||
|
r7622 | atexit.register(self.kill_bg_processes) | ||
|
r7618 | |||
def __del__(self): | ||||
self.kill_bg_processes() | ||||
|
r7299 | |||
def _generate_script_magics(self): | ||||
cell_magics = self.magics['cell'] | ||||
for name in self.script_magics: | ||||
cell_magics[name] = self._make_script_magic(name) | ||||
def _make_script_magic(self, name): | ||||
"""make a named magic, that calls %%script with a particular program""" | ||||
# expand to explicit path if necessary: | ||||
script = self.script_paths.get(name, name) | ||||
|
r7398 | @magic_arguments.magic_arguments() | ||
@script_args | ||||
|
r7299 | def named_script_magic(line, cell): | ||
# if line, add it as cl-flags | ||||
if line: | ||||
|
r26208 | line = "%s %s" % (script, line) | ||
|
r7299 | else: | ||
line = script | ||||
return self.shebang(line, cell) | ||||
# write a basic docstring: | ||||
named_script_magic.__doc__ = \ | ||||
"""%%{name} script magic | ||||
Run cells with {script} in a subprocess. | ||||
This is a shortcut for `%%script {script}` | ||||
""".format(**locals()) | ||||
return named_script_magic | ||||
|
r7398 | |||
@magic_arguments.magic_arguments() | ||||
@script_args | ||||
|
r7299 | @cell_magic("script") | ||
def shebang(self, line, cell): | ||||
"""Run a cell via a shell command | ||||
The `%%script` line is like the #! line of script, | ||||
specifying a program (bash, perl, ruby, etc.) with which to run. | ||||
The rest of the cell is run by that program. | ||||
Examples | ||||
-------- | ||||
:: | ||||
In [1]: %%script bash | ||||
...: for i in 1 2 3; do | ||||
...: echo $i | ||||
...: done | ||||
1 | ||||
2 | ||||
3 | ||||
""" | ||||
|
r26199 | |||
async def _handle_stream(stream, stream_arg, file_object): | ||||
while True: | ||||
line = (await stream.readline()).decode("utf8") | ||||
if not line: | ||||
break | ||||
if stream_arg: | ||||
self.shell.user_ns[stream_arg] = line | ||||
else: | ||||
file_object.write(line) | ||||
file_object.flush() | ||||
async def _stream_communicate(process, cell): | ||||
process.stdin.write(cell) | ||||
process.stdin.close() | ||||
stdout_task = asyncio.create_task( | ||||
_handle_stream(process.stdout, args.out, sys.stdout) | ||||
) | ||||
stderr_task = asyncio.create_task( | ||||
_handle_stream(process.stderr, args.err, sys.stderr) | ||||
) | ||||
await asyncio.wait([stdout_task, stderr_task]) | ||||
if sys.platform.startswith("win"): | ||||
asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) | ||||
loop = asyncio.get_event_loop() | ||||
|
r26200 | argv = arg_split(line, posix=not sys.platform.startswith("win")) | ||
|
r7398 | args, cmd = self.shebang.parser.parse_known_args(argv) | ||
|
r8470 | try: | ||
|
r26199 | p = loop.run_until_complete( | ||
asyncio.create_subprocess_exec( | ||||
*cmd, | ||||
stdout=asyncio.subprocess.PIPE, | ||||
stderr=asyncio.subprocess.PIPE, | ||||
stdin=asyncio.subprocess.PIPE | ||||
) | ||||
) | ||||
|
r8470 | except OSError as e: | ||
if e.errno == errno.ENOENT: | ||||
|
r13348 | print("Couldn't find program: %r" % cmd[0]) | ||
|
r8470 | return | ||
else: | ||||
raise | ||||
|
r7440 | |||
|
r17171 | if not cell.endswith('\n'): | ||
cell += '\n' | ||||
|
r7440 | cell = cell.encode('utf8', 'replace') | ||
|
r7398 | if args.bg: | ||
|
r7618 | self.bg_processes.append(p) | ||
|
r7630 | self._gc_bg_processes() | ||
|
r24375 | to_close = [] | ||
|
r7398 | if args.out: | ||
self.shell.user_ns[args.out] = p.stdout | ||||
|
r24375 | else: | ||
to_close.append(p.stdout) | ||||
|
r7398 | if args.err: | ||
|
r7407 | self.shell.user_ns[args.err] = p.stderr | ||
|
r24375 | else: | ||
to_close.append(p.stderr) | ||||
self.job_manager.new(self._run_script, p, cell, to_close, daemon=True) | ||||
|
r7549 | if args.proc: | ||
self.shell.user_ns[args.proc] = p | ||||
|
r7398 | return | ||
|
r7561 | try: | ||
|
r26199 | loop.run_until_complete(_stream_communicate(p, cell)) | ||
|
r7561 | except KeyboardInterrupt: | ||
|
r7562 | try: | ||
p.send_signal(signal.SIGINT) | ||||
time.sleep(0.1) | ||||
|
r26199 | if p.returncode is not None: | ||
|
r13348 | print("Process is interrupted.") | ||
|
r7562 | return | ||
p.terminate() | ||||
time.sleep(0.1) | ||||
|
r26199 | if p.returncode is not None: | ||
|
r13348 | print("Process is terminated.") | ||
|
r7562 | return | ||
p.kill() | ||||
|
r13348 | print("Process is killed.") | ||
|
r7565 | except OSError: | ||
pass | ||||
except Exception as e: | ||||
|
r26208 | print("Error while terminating subprocess (pid=%i): %s" % (p.pid, e)) | ||
|
r7561 | return | ||
|
r24410 | if args.raise_error and p.returncode!=0: | ||
|
r26199 | raise CalledProcessError(p.returncode, cell) | ||
|
r7299 | |||
|
r24375 | def _run_script(self, p, cell, to_close): | ||
|
r7398 | """callback for running the script in the background""" | ||
p.stdin.write(cell) | ||||
p.stdin.close() | ||||
|
r24375 | for s in to_close: | ||
s.close() | ||||
|
r7398 | p.wait() | ||
|
r7618 | |||
|
r7619 | @line_magic("killbgscripts") | ||
|
r7629 | def killbgscripts(self, _nouse_=''): | ||
"""Kill all BG processes started by %%script and its family.""" | ||||
self.kill_bg_processes() | ||||
|
r13348 | print("All background processes were killed.") | ||
|
r7629 | |||
def kill_bg_processes(self): | ||||
|
r7618 | """Kill all BG processes which are still running.""" | ||
|
r23508 | if not self.bg_processes: | ||
return | ||||
|
r7618 | for p in self.bg_processes: | ||
|
r26199 | if p.returncode is None: | ||
|
r7618 | try: | ||
p.send_signal(signal.SIGINT) | ||||
except: | ||||
pass | ||||
time.sleep(0.1) | ||||
|
r23508 | self._gc_bg_processes() | ||
if not self.bg_processes: | ||||
return | ||||
|
r7618 | for p in self.bg_processes: | ||
|
r26199 | if p.returncode is None: | ||
|
r7618 | try: | ||
p.terminate() | ||||
except: | ||||
pass | ||||
time.sleep(0.1) | ||||
|
r23508 | self._gc_bg_processes() | ||
if not self.bg_processes: | ||||
return | ||||
|
r7618 | for p in self.bg_processes: | ||
|
r26199 | if p.returncode is None: | ||
|
r7618 | try: | ||
p.kill() | ||||
except: | ||||
pass | ||||
|
r7630 | self._gc_bg_processes() | ||
def _gc_bg_processes(self): | ||||
|
r26199 | self.bg_processes = [p for p in self.bg_processes if p.returncode is None] | ||