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