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