script.py
279 lines
| 8.6 KiB
| text/x-python
|
PythonLexer
MinRK
|
r7299 | """Magic functions for running cells in various scripts.""" | ||
Min RK
|
r22340 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
MinRK
|
r7299 | |||
MinRK
|
r8470 | import errno | ||
MinRK
|
r7299 | import os | ||
import sys | ||||
Takafumi Arakaki
|
r7561 | import signal | ||
import time | ||||
MinRK
|
r7299 | from subprocess import Popen, PIPE | ||
Takafumi Arakaki
|
r7622 | import atexit | ||
MinRK
|
r7299 | |||
MinRK
|
r7398 | from IPython.core import magic_arguments | ||
MinRK
|
r7299 | from IPython.core.magic import ( | ||
Magics, magics_class, line_magic, cell_magic | ||||
) | ||||
MinRK
|
r7398 | from IPython.lib.backgroundjobs import BackgroundJobManager | ||
from IPython.utils import py3compat | ||||
MinRK
|
r8469 | from IPython.utils.process import arg_split | ||
Min RK
|
r22340 | from traitlets import List, Dict, default | ||
MinRK
|
r7299 | |||
#----------------------------------------------------------------------------- | ||||
# Magic implementation classes | ||||
#----------------------------------------------------------------------------- | ||||
MinRK
|
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*, | ||||
instead of the stderr text itself. | ||||
""" | ||||
), | ||||
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*, | ||||
instead of the stderr text itself. | ||||
""" | ||||
), | ||||
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. | ||||
""" | ||||
), | ||||
Takafumi Arakaki
|
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. | ||||
""" | ||||
), | ||||
MinRK
|
r7398 | ] | ||
for arg in args: | ||||
f = arg(f) | ||||
return f | ||||
MinRK
|
r7299 | @magics_class | ||
Matthias BUSSONNIER
|
r13237 | class ScriptMagics(Magics): | ||
MinRK
|
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. | ||||
""" | ||||
Min RK
|
r22340 | script_magics = List( | ||
MinRK
|
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 | ||||
""", | ||||
Min RK
|
r22340 | ).tag(config=True) | ||
@default('script_magics') | ||||
MinRK
|
r7299 | def _script_magics_default(self): | ||
MinRK
|
r8469 | """default to a common list of programs""" | ||
MinRK
|
r7299 | |||
MinRK
|
r8469 | defaults = [ | ||
MinRK
|
r7299 | 'sh', | ||
'bash', | ||||
'perl', | ||||
'ruby', | ||||
MinRK
|
r8469 | 'python', | ||
Thomas Kluyver
|
r16020 | 'python2', | ||
MinRK
|
r7299 | 'python3', | ||
'pypy', | ||||
MinRK
|
r8469 | ] | ||
if os.name == 'nt': | ||||
defaults.extend([ | ||||
'cmd', | ||||
]) | ||||
MinRK
|
r7299 | |||
return defaults | ||||
Min RK
|
r22340 | script_paths = Dict( | ||
MinRK
|
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. | ||||
""" | ||||
Min RK
|
r22340 | ).tag(config=True) | ||
MinRK
|
r7299 | |||
def __init__(self, shell=None): | ||||
Matthias BUSSONNIER
|
r13237 | super(ScriptMagics, self).__init__(shell=shell) | ||
MinRK
|
r7299 | self._generate_script_magics() | ||
MinRK
|
r7398 | self.job_manager = BackgroundJobManager() | ||
Takafumi Arakaki
|
r7618 | self.bg_processes = [] | ||
Takafumi Arakaki
|
r7622 | atexit.register(self.kill_bg_processes) | ||
Takafumi Arakaki
|
r7618 | |||
def __del__(self): | ||||
self.kill_bg_processes() | ||||
MinRK
|
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) | ||||
MinRK
|
r7398 | @magic_arguments.magic_arguments() | ||
@script_args | ||||
MinRK
|
r7299 | def named_script_magic(line, cell): | ||
# if line, add it as cl-flags | ||||
if line: | ||||
line = "%s %s" % (script, line) | ||||
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 | ||||
MinRK
|
r7398 | |||
@magic_arguments.magic_arguments() | ||||
@script_args | ||||
MinRK
|
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 | ||||
""" | ||||
MinRK
|
r7398 | argv = arg_split(line, posix = not sys.platform.startswith('win')) | ||
args, cmd = self.shebang.parser.parse_known_args(argv) | ||||
MinRK
|
r8470 | |||
try: | ||||
p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE) | ||||
except OSError as e: | ||||
if e.errno == errno.ENOENT: | ||||
Thomas Kluyver
|
r13348 | print("Couldn't find program: %r" % cmd[0]) | ||
MinRK
|
r8470 | return | ||
else: | ||||
raise | ||||
MinRK
|
r7440 | |||
Eric Galloway
|
r17171 | if not cell.endswith('\n'): | ||
cell += '\n' | ||||
MinRK
|
r7440 | cell = cell.encode('utf8', 'replace') | ||
MinRK
|
r7398 | if args.bg: | ||
Takafumi Arakaki
|
r7618 | self.bg_processes.append(p) | ||
Takafumi Arakaki
|
r7630 | self._gc_bg_processes() | ||
MinRK
|
r7398 | if args.out: | ||
self.shell.user_ns[args.out] = p.stdout | ||||
if args.err: | ||||
MinRK
|
r7407 | self.shell.user_ns[args.err] = p.stderr | ||
Takafumi Arakaki
|
r7621 | self.job_manager.new(self._run_script, p, cell, daemon=True) | ||
Takafumi Arakaki
|
r7549 | if args.proc: | ||
self.shell.user_ns[args.proc] = p | ||||
MinRK
|
r7398 | return | ||
Takafumi Arakaki
|
r7561 | try: | ||
out, err = p.communicate(cell) | ||||
except KeyboardInterrupt: | ||||
Takafumi Arakaki
|
r7562 | try: | ||
p.send_signal(signal.SIGINT) | ||||
time.sleep(0.1) | ||||
if p.poll() is not None: | ||||
Thomas Kluyver
|
r13348 | print("Process is interrupted.") | ||
Takafumi Arakaki
|
r7562 | return | ||
p.terminate() | ||||
time.sleep(0.1) | ||||
if p.poll() is not None: | ||||
Thomas Kluyver
|
r13348 | print("Process is terminated.") | ||
Takafumi Arakaki
|
r7562 | return | ||
p.kill() | ||||
Thomas Kluyver
|
r13348 | print("Process is killed.") | ||
Takafumi Arakaki
|
r7565 | except OSError: | ||
pass | ||||
except Exception as e: | ||||
Thomas Kluyver
|
r13348 | print("Error while terminating subprocess (pid=%i): %s" \ | ||
% (p.pid, e)) | ||||
Takafumi Arakaki
|
r7561 | return | ||
MinRK
|
r7398 | out = py3compat.bytes_to_str(out) | ||
err = py3compat.bytes_to_str(err) | ||||
if args.out: | ||||
self.shell.user_ns[args.out] = out | ||||
else: | ||||
sys.stdout.write(out) | ||||
sys.stdout.flush() | ||||
if args.err: | ||||
self.shell.user_ns[args.err] = err | ||||
else: | ||||
sys.stderr.write(err) | ||||
sys.stderr.flush() | ||||
MinRK
|
r7299 | |||
MinRK
|
r7398 | def _run_script(self, p, cell): | ||
"""callback for running the script in the background""" | ||||
p.stdin.write(cell) | ||||
p.stdin.close() | ||||
p.wait() | ||||
Takafumi Arakaki
|
r7618 | |||
Takafumi Arakaki
|
r7619 | @line_magic("killbgscripts") | ||
Takafumi Arakaki
|
r7629 | def killbgscripts(self, _nouse_=''): | ||
"""Kill all BG processes started by %%script and its family.""" | ||||
self.kill_bg_processes() | ||||
Thomas Kluyver
|
r13348 | print("All background processes were killed.") | ||
Takafumi Arakaki
|
r7629 | |||
def kill_bg_processes(self): | ||||
Takafumi Arakaki
|
r7618 | """Kill all BG processes which are still running.""" | ||
Leo Singer
|
r23508 | if not self.bg_processes: | ||
return | ||||
Takafumi Arakaki
|
r7618 | for p in self.bg_processes: | ||
if p.poll() is None: | ||||
try: | ||||
p.send_signal(signal.SIGINT) | ||||
except: | ||||
pass | ||||
time.sleep(0.1) | ||||
Leo Singer
|
r23508 | self._gc_bg_processes() | ||
if not self.bg_processes: | ||||
return | ||||
Takafumi Arakaki
|
r7618 | for p in self.bg_processes: | ||
if p.poll() is None: | ||||
try: | ||||
p.terminate() | ||||
except: | ||||
pass | ||||
time.sleep(0.1) | ||||
Leo Singer
|
r23508 | self._gc_bg_processes() | ||
if not self.bg_processes: | ||||
return | ||||
Takafumi Arakaki
|
r7618 | for p in self.bg_processes: | ||
if p.poll() is None: | ||||
try: | ||||
p.kill() | ||||
except: | ||||
pass | ||||
Takafumi Arakaki
|
r7630 | self._gc_bg_processes() | ||
def _gc_bg_processes(self): | ||||
self.bg_processes = [p for p in self.bg_processes if p.poll() is None] | ||||