##// END OF EJS Templates
Merge pull request #2424 from minrk/fastscript...
Bussonnier Matthias -
r8504:cfe1161c merge
parent child Browse files
Show More
@@ -1,288 +1,283 b''
1 """Magic functions for running cells in various scripts."""
1 """Magic functions for running cells in various scripts."""
2 #-----------------------------------------------------------------------------
2 #-----------------------------------------------------------------------------
3 # Copyright (c) 2012 The IPython Development Team.
3 # Copyright (c) 2012 The IPython Development Team.
4 #
4 #
5 # Distributed under the terms of the Modified BSD License.
5 # Distributed under the terms of the Modified BSD License.
6 #
6 #
7 # The full license is in the file COPYING.txt, distributed with this software.
7 # The full license is in the file COPYING.txt, distributed with this software.
8 #-----------------------------------------------------------------------------
8 #-----------------------------------------------------------------------------
9
9
10 #-----------------------------------------------------------------------------
10 #-----------------------------------------------------------------------------
11 # Imports
11 # Imports
12 #-----------------------------------------------------------------------------
12 #-----------------------------------------------------------------------------
13
13
14 # Stdlib
14 # Stdlib
15 import errno
15 import os
16 import os
16 import re
17 import re
17 import sys
18 import sys
18 import signal
19 import signal
19 import time
20 import time
20 from subprocess import Popen, PIPE
21 from subprocess import Popen, PIPE
21 import atexit
22 import atexit
22
23
23 # Our own packages
24 # Our own packages
24 from IPython.config.configurable import Configurable
25 from IPython.config.configurable import Configurable
25 from IPython.core import magic_arguments
26 from IPython.core import magic_arguments
26 from IPython.core.error import UsageError
27 from IPython.core.error import UsageError
27 from IPython.core.magic import (
28 from IPython.core.magic import (
28 Magics, magics_class, line_magic, cell_magic
29 Magics, magics_class, line_magic, cell_magic
29 )
30 )
30 from IPython.lib.backgroundjobs import BackgroundJobManager
31 from IPython.lib.backgroundjobs import BackgroundJobManager
31 from IPython.testing.skipdoctest import skip_doctest
32 from IPython.testing.skipdoctest import skip_doctest
32 from IPython.utils import py3compat
33 from IPython.utils import py3compat
33 from IPython.utils.process import find_cmd, FindCmdError, arg_split
34 from IPython.utils.process import arg_split
34 from IPython.utils.traitlets import List, Dict
35 from IPython.utils.traitlets import List, Dict
35
36
36 #-----------------------------------------------------------------------------
37 #-----------------------------------------------------------------------------
37 # Magic implementation classes
38 # Magic implementation classes
38 #-----------------------------------------------------------------------------
39 #-----------------------------------------------------------------------------
39
40
40 def script_args(f):
41 def script_args(f):
41 """single decorator for adding script args"""
42 """single decorator for adding script args"""
42 args = [
43 args = [
43 magic_arguments.argument(
44 magic_arguments.argument(
44 '--out', type=str,
45 '--out', type=str,
45 help="""The variable in which to store stdout from the script.
46 help="""The variable in which to store stdout from the script.
46 If the script is backgrounded, this will be the stdout *pipe*,
47 If the script is backgrounded, this will be the stdout *pipe*,
47 instead of the stderr text itself.
48 instead of the stderr text itself.
48 """
49 """
49 ),
50 ),
50 magic_arguments.argument(
51 magic_arguments.argument(
51 '--err', type=str,
52 '--err', type=str,
52 help="""The variable in which to store stderr from the script.
53 help="""The variable in which to store stderr from the script.
53 If the script is backgrounded, this will be the stderr *pipe*,
54 If the script is backgrounded, this will be the stderr *pipe*,
54 instead of the stderr text itself.
55 instead of the stderr text itself.
55 """
56 """
56 ),
57 ),
57 magic_arguments.argument(
58 magic_arguments.argument(
58 '--bg', action="store_true",
59 '--bg', action="store_true",
59 help="""Whether to run the script in the background.
60 help="""Whether to run the script in the background.
60 If given, the only way to see the output of the command is
61 If given, the only way to see the output of the command is
61 with --out/err.
62 with --out/err.
62 """
63 """
63 ),
64 ),
64 magic_arguments.argument(
65 magic_arguments.argument(
65 '--proc', type=str,
66 '--proc', type=str,
66 help="""The variable in which to store Popen instance.
67 help="""The variable in which to store Popen instance.
67 This is used only when --bg option is given.
68 This is used only when --bg option is given.
68 """
69 """
69 ),
70 ),
70 ]
71 ]
71 for arg in args:
72 for arg in args:
72 f = arg(f)
73 f = arg(f)
73 return f
74 return f
74
75
75 @magics_class
76 @magics_class
76 class ScriptMagics(Magics, Configurable):
77 class ScriptMagics(Magics, Configurable):
77 """Magics for talking to scripts
78 """Magics for talking to scripts
78
79
79 This defines a base `%%script` cell magic for running a cell
80 This defines a base `%%script` cell magic for running a cell
80 with a program in a subprocess, and registers a few top-level
81 with a program in a subprocess, and registers a few top-level
81 magics that call %%script with common interpreters.
82 magics that call %%script with common interpreters.
82 """
83 """
83 script_magics = List(config=True,
84 script_magics = List(config=True,
84 help="""Extra script cell magics to define
85 help="""Extra script cell magics to define
85
86
86 This generates simple wrappers of `%%script foo` as `%%foo`.
87 This generates simple wrappers of `%%script foo` as `%%foo`.
87
88
88 If you want to add script magics that aren't on your path,
89 If you want to add script magics that aren't on your path,
89 specify them in script_paths
90 specify them in script_paths
90 """,
91 """,
91 )
92 )
92 def _script_magics_default(self):
93 def _script_magics_default(self):
93 """default to a common list of programs if we find them"""
94 """default to a common list of programs"""
94
95
95 defaults = []
96 defaults = [
96 to_try = []
97 if os.name == 'nt':
98 defaults.append('cmd')
99 to_try.append('powershell')
100 to_try.extend([
101 'sh',
97 'sh',
102 'bash',
98 'bash',
103 'perl',
99 'perl',
104 'ruby',
100 'ruby',
101 'python',
105 'python3',
102 'python3',
106 'pypy',
103 'pypy',
107 ])
104 ]
105 if os.name == 'nt':
106 defaults.extend([
107 'cmd',
108 'powershell',
109 ])
108
110
109 for cmd in to_try:
110 if cmd in self.script_paths:
111 defaults.append(cmd)
112 else:
113 try:
114 find_cmd(cmd)
115 except FindCmdError:
116 # command not found, ignore it
117 pass
118 except ImportError:
119 # Windows without pywin32, find_cmd doesn't work
120 pass
121 else:
122 defaults.append(cmd)
123 return defaults
111 return defaults
124
112
125 script_paths = Dict(config=True,
113 script_paths = Dict(config=True,
126 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
114 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
127
115
128 Only necessary for items in script_magics where the default path will not
116 Only necessary for items in script_magics where the default path will not
129 find the right interpreter.
117 find the right interpreter.
130 """
118 """
131 )
119 )
132
120
133 def __init__(self, shell=None):
121 def __init__(self, shell=None):
134 Configurable.__init__(self, config=shell.config)
122 Configurable.__init__(self, config=shell.config)
135 self._generate_script_magics()
123 self._generate_script_magics()
136 Magics.__init__(self, shell=shell)
124 Magics.__init__(self, shell=shell)
137 self.job_manager = BackgroundJobManager()
125 self.job_manager = BackgroundJobManager()
138 self.bg_processes = []
126 self.bg_processes = []
139 atexit.register(self.kill_bg_processes)
127 atexit.register(self.kill_bg_processes)
140
128
141 def __del__(self):
129 def __del__(self):
142 self.kill_bg_processes()
130 self.kill_bg_processes()
143
131
144 def _generate_script_magics(self):
132 def _generate_script_magics(self):
145 cell_magics = self.magics['cell']
133 cell_magics = self.magics['cell']
146 for name in self.script_magics:
134 for name in self.script_magics:
147 cell_magics[name] = self._make_script_magic(name)
135 cell_magics[name] = self._make_script_magic(name)
148
136
149 def _make_script_magic(self, name):
137 def _make_script_magic(self, name):
150 """make a named magic, that calls %%script with a particular program"""
138 """make a named magic, that calls %%script with a particular program"""
151 # expand to explicit path if necessary:
139 # expand to explicit path if necessary:
152 script = self.script_paths.get(name, name)
140 script = self.script_paths.get(name, name)
153
141
154 @magic_arguments.magic_arguments()
142 @magic_arguments.magic_arguments()
155 @script_args
143 @script_args
156 def named_script_magic(line, cell):
144 def named_script_magic(line, cell):
157 # if line, add it as cl-flags
145 # if line, add it as cl-flags
158 if line:
146 if line:
159 line = "%s %s" % (script, line)
147 line = "%s %s" % (script, line)
160 else:
148 else:
161 line = script
149 line = script
162 return self.shebang(line, cell)
150 return self.shebang(line, cell)
163
151
164 # write a basic docstring:
152 # write a basic docstring:
165 named_script_magic.__doc__ = \
153 named_script_magic.__doc__ = \
166 """%%{name} script magic
154 """%%{name} script magic
167
155
168 Run cells with {script} in a subprocess.
156 Run cells with {script} in a subprocess.
169
157
170 This is a shortcut for `%%script {script}`
158 This is a shortcut for `%%script {script}`
171 """.format(**locals())
159 """.format(**locals())
172
160
173 return named_script_magic
161 return named_script_magic
174
162
175 @magic_arguments.magic_arguments()
163 @magic_arguments.magic_arguments()
176 @script_args
164 @script_args
177 @cell_magic("script")
165 @cell_magic("script")
178 def shebang(self, line, cell):
166 def shebang(self, line, cell):
179 """Run a cell via a shell command
167 """Run a cell via a shell command
180
168
181 The `%%script` line is like the #! line of script,
169 The `%%script` line is like the #! line of script,
182 specifying a program (bash, perl, ruby, etc.) with which to run.
170 specifying a program (bash, perl, ruby, etc.) with which to run.
183
171
184 The rest of the cell is run by that program.
172 The rest of the cell is run by that program.
185
173
186 Examples
174 Examples
187 --------
175 --------
188 ::
176 ::
189
177
190 In [1]: %%script bash
178 In [1]: %%script bash
191 ...: for i in 1 2 3; do
179 ...: for i in 1 2 3; do
192 ...: echo $i
180 ...: echo $i
193 ...: done
181 ...: done
194 1
182 1
195 2
183 2
196 3
184 3
197 """
185 """
198 argv = arg_split(line, posix = not sys.platform.startswith('win'))
186 argv = arg_split(line, posix = not sys.platform.startswith('win'))
199 args, cmd = self.shebang.parser.parse_known_args(argv)
187 args, cmd = self.shebang.parser.parse_known_args(argv)
200
188
201 p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE)
189 try:
190 p = Popen(cmd, stdout=PIPE, stderr=PIPE, stdin=PIPE)
191 except OSError as e:
192 if e.errno == errno.ENOENT:
193 print "Couldn't find program: %r" % cmd[0]
194 return
195 else:
196 raise
202
197
203 cell = cell.encode('utf8', 'replace')
198 cell = cell.encode('utf8', 'replace')
204 if args.bg:
199 if args.bg:
205 self.bg_processes.append(p)
200 self.bg_processes.append(p)
206 self._gc_bg_processes()
201 self._gc_bg_processes()
207 if args.out:
202 if args.out:
208 self.shell.user_ns[args.out] = p.stdout
203 self.shell.user_ns[args.out] = p.stdout
209 if args.err:
204 if args.err:
210 self.shell.user_ns[args.err] = p.stderr
205 self.shell.user_ns[args.err] = p.stderr
211 self.job_manager.new(self._run_script, p, cell, daemon=True)
206 self.job_manager.new(self._run_script, p, cell, daemon=True)
212 if args.proc:
207 if args.proc:
213 self.shell.user_ns[args.proc] = p
208 self.shell.user_ns[args.proc] = p
214 return
209 return
215
210
216 try:
211 try:
217 out, err = p.communicate(cell)
212 out, err = p.communicate(cell)
218 except KeyboardInterrupt:
213 except KeyboardInterrupt:
219 try:
214 try:
220 p.send_signal(signal.SIGINT)
215 p.send_signal(signal.SIGINT)
221 time.sleep(0.1)
216 time.sleep(0.1)
222 if p.poll() is not None:
217 if p.poll() is not None:
223 print "Process is interrupted."
218 print "Process is interrupted."
224 return
219 return
225 p.terminate()
220 p.terminate()
226 time.sleep(0.1)
221 time.sleep(0.1)
227 if p.poll() is not None:
222 if p.poll() is not None:
228 print "Process is terminated."
223 print "Process is terminated."
229 return
224 return
230 p.kill()
225 p.kill()
231 print "Process is killed."
226 print "Process is killed."
232 except OSError:
227 except OSError:
233 pass
228 pass
234 except Exception as e:
229 except Exception as e:
235 print "Error while terminating subprocess (pid=%i): %s" \
230 print "Error while terminating subprocess (pid=%i): %s" \
236 % (p.pid, e)
231 % (p.pid, e)
237 return
232 return
238 out = py3compat.bytes_to_str(out)
233 out = py3compat.bytes_to_str(out)
239 err = py3compat.bytes_to_str(err)
234 err = py3compat.bytes_to_str(err)
240 if args.out:
235 if args.out:
241 self.shell.user_ns[args.out] = out
236 self.shell.user_ns[args.out] = out
242 else:
237 else:
243 sys.stdout.write(out)
238 sys.stdout.write(out)
244 sys.stdout.flush()
239 sys.stdout.flush()
245 if args.err:
240 if args.err:
246 self.shell.user_ns[args.err] = err
241 self.shell.user_ns[args.err] = err
247 else:
242 else:
248 sys.stderr.write(err)
243 sys.stderr.write(err)
249 sys.stderr.flush()
244 sys.stderr.flush()
250
245
251 def _run_script(self, p, cell):
246 def _run_script(self, p, cell):
252 """callback for running the script in the background"""
247 """callback for running the script in the background"""
253 p.stdin.write(cell)
248 p.stdin.write(cell)
254 p.stdin.close()
249 p.stdin.close()
255 p.wait()
250 p.wait()
256
251
257 @line_magic("killbgscripts")
252 @line_magic("killbgscripts")
258 def killbgscripts(self, _nouse_=''):
253 def killbgscripts(self, _nouse_=''):
259 """Kill all BG processes started by %%script and its family."""
254 """Kill all BG processes started by %%script and its family."""
260 self.kill_bg_processes()
255 self.kill_bg_processes()
261 print "All background processes were killed."
256 print "All background processes were killed."
262
257
263 def kill_bg_processes(self):
258 def kill_bg_processes(self):
264 """Kill all BG processes which are still running."""
259 """Kill all BG processes which are still running."""
265 for p in self.bg_processes:
260 for p in self.bg_processes:
266 if p.poll() is None:
261 if p.poll() is None:
267 try:
262 try:
268 p.send_signal(signal.SIGINT)
263 p.send_signal(signal.SIGINT)
269 except:
264 except:
270 pass
265 pass
271 time.sleep(0.1)
266 time.sleep(0.1)
272 for p in self.bg_processes:
267 for p in self.bg_processes:
273 if p.poll() is None:
268 if p.poll() is None:
274 try:
269 try:
275 p.terminate()
270 p.terminate()
276 except:
271 except:
277 pass
272 pass
278 time.sleep(0.1)
273 time.sleep(0.1)
279 for p in self.bg_processes:
274 for p in self.bg_processes:
280 if p.poll() is None:
275 if p.poll() is None:
281 try:
276 try:
282 p.kill()
277 p.kill()
283 except:
278 except:
284 pass
279 pass
285 self._gc_bg_processes()
280 self._gc_bg_processes()
286
281
287 def _gc_bg_processes(self):
282 def _gc_bg_processes(self):
288 self.bg_processes = [p for p in self.bg_processes if p.poll() is None]
283 self.bg_processes = [p for p in self.bg_processes if p.poll() is None]
General Comments 0
You need to be logged in to leave comments. Login now