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