##// END OF EJS Templates
Update script.py...
Garland Zhang -
Show More
@@ -1,362 +1,362 b''
1 1 """Magic functions for running cells in various scripts."""
2 2
3 3 # Copyright (c) IPython Development Team.
4 4 # Distributed under the terms of the Modified BSD License.
5 5
6 6 import asyncio
7 7 import atexit
8 8 import errno
9 9 import os
10 10 import signal
11 11 import sys
12 12 import time
13 13 from subprocess import CalledProcessError
14 14 from threading import Thread
15 15
16 16 from traitlets import Any, Dict, List, default
17 17
18 18 from IPython.core import magic_arguments
19 19 from IPython.core.async_helpers import _AsyncIOProxy
20 20 from IPython.core.magic import Magics, cell_magic, line_magic, magics_class
21 21 from IPython.utils.process import arg_split
22 22
23 23 #-----------------------------------------------------------------------------
24 24 # Magic implementation classes
25 25 #-----------------------------------------------------------------------------
26 26
27 27 def script_args(f):
28 28 """single decorator for adding script args"""
29 29 args = [
30 30 magic_arguments.argument(
31 31 '--out', type=str,
32 32 help="""The variable in which to store stdout from the script.
33 33 If the script is backgrounded, this will be the stdout *pipe*,
34 34 instead of the stderr text itself and will not be auto closed.
35 35 """
36 36 ),
37 37 magic_arguments.argument(
38 38 '--err', type=str,
39 39 help="""The variable in which to store stderr from the script.
40 40 If the script is backgrounded, this will be the stderr *pipe*,
41 41 instead of the stderr text itself and will not be autoclosed.
42 42 """
43 43 ),
44 44 magic_arguments.argument(
45 45 '--bg', action="store_true",
46 46 help="""Whether to run the script in the background.
47 47 If given, the only way to see the output of the command is
48 48 with --out/err.
49 49 """
50 50 ),
51 51 magic_arguments.argument(
52 52 '--proc', type=str,
53 53 help="""The variable in which to store Popen instance.
54 54 This is used only when --bg option is given.
55 55 """
56 56 ),
57 57 magic_arguments.argument(
58 58 '--no-raise-error', action="store_false", dest='raise_error',
59 59 help="""Whether you should raise an error message in addition to
60 60 a stream on stderr if you get a nonzero exit code.
61 61 """,
62 62 ),
63 63 ]
64 64 for arg in args:
65 65 f = arg(f)
66 66 return f
67 67
68 68
69 69 @magics_class
70 70 class ScriptMagics(Magics):
71 71 """Magics for talking to scripts
72 72
73 73 This defines a base `%%script` cell magic for running a cell
74 74 with a program in a subprocess, and registers a few top-level
75 75 magics that call %%script with common interpreters.
76 76 """
77 77
78 78 event_loop = Any(
79 79 help="""
80 80 The event loop on which to run subprocesses
81 81
82 82 Not the main event loop,
83 83 because we want to be able to make blocking calls
84 84 and have certain requirements we don't want to impose on the main loop.
85 85 """
86 86 )
87 87
88 88 script_magics = List(
89 89 help="""Extra script cell magics to define
90 90
91 91 This generates simple wrappers of `%%script foo` as `%%foo`.
92 92
93 93 If you want to add script magics that aren't on your path,
94 94 specify them in script_paths
95 95 """,
96 96 ).tag(config=True)
97 97 @default('script_magics')
98 98 def _script_magics_default(self):
99 99 """default to a common list of programs"""
100 100
101 101 defaults = [
102 102 'sh',
103 103 'bash',
104 104 'perl',
105 105 'ruby',
106 106 'python',
107 107 'python2',
108 108 'python3',
109 109 'pypy',
110 110 ]
111 111 if os.name == 'nt':
112 112 defaults.extend([
113 113 'cmd',
114 114 ])
115 115
116 116 return defaults
117 117
118 118 script_paths = Dict(
119 119 help="""Dict mapping short 'ruby' names to full paths, such as '/opt/secret/bin/ruby'
120 120
121 121 Only necessary for items in script_magics where the default path will not
122 122 find the right interpreter.
123 123 """
124 124 ).tag(config=True)
125 125
126 126 def __init__(self, shell=None):
127 127 super(ScriptMagics, self).__init__(shell=shell)
128 128 self._generate_script_magics()
129 129 self.bg_processes = []
130 130 atexit.register(self.kill_bg_processes)
131 131
132 132 def __del__(self):
133 133 self.kill_bg_processes()
134 134
135 135 def _generate_script_magics(self):
136 136 cell_magics = self.magics['cell']
137 137 for name in self.script_magics:
138 138 cell_magics[name] = self._make_script_magic(name)
139 139
140 140 def _make_script_magic(self, name):
141 141 """make a named magic, that calls %%script with a particular program"""
142 142 # expand to explicit path if necessary:
143 143 script = self.script_paths.get(name, name)
144 144
145 145 @magic_arguments.magic_arguments()
146 146 @script_args
147 147 def named_script_magic(line, cell):
148 148 # if line, add it as cl-flags
149 149 if line:
150 150 line = "%s %s" % (script, line)
151 151 else:
152 152 line = script
153 153 return self.shebang(line, cell)
154 154
155 155 # write a basic docstring:
156 156 named_script_magic.__doc__ = \
157 157 """%%{name} script magic
158 158
159 159 Run cells with {script} in a subprocess.
160 160
161 161 This is a shortcut for `%%script {script}`
162 162 """.format(**locals())
163 163
164 164 return named_script_magic
165 165
166 166 @magic_arguments.magic_arguments()
167 167 @script_args
168 168 @cell_magic("script")
169 169 def shebang(self, line, cell):
170 170 """Run a cell via a shell command
171 171
172 172 The `%%script` line is like the #! line of script,
173 173 specifying a program (bash, perl, ruby, etc.) with which to run.
174 174
175 175 The rest of the cell is run by that program.
176 176
177 177 Examples
178 178 --------
179 179 ::
180 180
181 181 In [1]: %%script bash
182 182 ...: for i in 1 2 3; do
183 183 ...: echo $i
184 184 ...: done
185 185 1
186 186 2
187 187 3
188 188 """
189 189
190 190 # Create the event loop in which to run script magics
191 191 # this operates on a background thread
192 192 if self.event_loop is None:
193 193 if sys.platform == "win32":
194 194 # don't override the current policy,
195 195 # just create an event loop
196 196 event_loop = asyncio.WindowsProactorEventLoopPolicy().new_event_loop()
197 197 else:
198 198 event_loop = asyncio.new_event_loop()
199 199 self.event_loop = event_loop
200 200
201 201 # start the loop in a background thread
202 202 asyncio_thread = Thread(target=event_loop.run_forever, daemon=True)
203 203 asyncio_thread.start()
204 204 else:
205 205 event_loop = self.event_loop
206 206
207 207 def in_thread(coro):
208 208 """Call a coroutine on the asyncio thread"""
209 209 return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
210 210
211 211 async def _handle_stream(stream, stream_arg, file_object):
212 212 while True:
213 line = (await stream.readline()).decode("utf8")
213 line = (await stream.readline()).decode("utf8", errors="replace")
214 214 if not line:
215 215 break
216 216 if stream_arg:
217 217 self.shell.user_ns[stream_arg] = line
218 218 else:
219 219 file_object.write(line)
220 220 file_object.flush()
221 221
222 222 async def _stream_communicate(process, cell):
223 223 process.stdin.write(cell)
224 224 process.stdin.close()
225 225 stdout_task = asyncio.create_task(
226 226 _handle_stream(process.stdout, args.out, sys.stdout)
227 227 )
228 228 stderr_task = asyncio.create_task(
229 229 _handle_stream(process.stderr, args.err, sys.stderr)
230 230 )
231 231 await asyncio.wait([stdout_task, stderr_task])
232 232 await process.wait()
233 233
234 234 argv = arg_split(line, posix=not sys.platform.startswith("win"))
235 235 args, cmd = self.shebang.parser.parse_known_args(argv)
236 236
237 237 try:
238 238 p = in_thread(
239 239 asyncio.create_subprocess_exec(
240 240 *cmd,
241 241 stdout=asyncio.subprocess.PIPE,
242 242 stderr=asyncio.subprocess.PIPE,
243 243 stdin=asyncio.subprocess.PIPE,
244 244 )
245 245 )
246 246 except OSError as e:
247 247 if e.errno == errno.ENOENT:
248 248 print("Couldn't find program: %r" % cmd[0])
249 249 return
250 250 else:
251 251 raise
252 252
253 253 if not cell.endswith('\n'):
254 254 cell += '\n'
255 255 cell = cell.encode('utf8', 'replace')
256 256 if args.bg:
257 257 self.bg_processes.append(p)
258 258 self._gc_bg_processes()
259 259 to_close = []
260 260 if args.out:
261 261 self.shell.user_ns[args.out] = _AsyncIOProxy(p.stdout, event_loop)
262 262 else:
263 263 to_close.append(p.stdout)
264 264 if args.err:
265 265 self.shell.user_ns[args.err] = _AsyncIOProxy(p.stderr, event_loop)
266 266 else:
267 267 to_close.append(p.stderr)
268 268 event_loop.call_soon_threadsafe(
269 269 lambda: asyncio.Task(self._run_script(p, cell, to_close))
270 270 )
271 271 if args.proc:
272 272 proc_proxy = _AsyncIOProxy(p, event_loop)
273 273 proc_proxy.stdout = _AsyncIOProxy(p.stdout, event_loop)
274 274 proc_proxy.stderr = _AsyncIOProxy(p.stderr, event_loop)
275 275 self.shell.user_ns[args.proc] = proc_proxy
276 276 return
277 277
278 278 try:
279 279 in_thread(_stream_communicate(p, cell))
280 280 except KeyboardInterrupt:
281 281 try:
282 282 p.send_signal(signal.SIGINT)
283 283 in_thread(asyncio.wait_for(p.wait(), timeout=0.1))
284 284 if p.returncode is not None:
285 285 print("Process is interrupted.")
286 286 return
287 287 p.terminate()
288 288 in_thread(asyncio.wait_for(p.wait(), timeout=0.1))
289 289 if p.returncode is not None:
290 290 print("Process is terminated.")
291 291 return
292 292 p.kill()
293 293 print("Process is killed.")
294 294 except OSError:
295 295 pass
296 296 except Exception as e:
297 297 print("Error while terminating subprocess (pid=%i): %s" % (p.pid, e))
298 298 return
299 299
300 300 if args.raise_error and p.returncode != 0:
301 301 # If we get here and p.returncode is still None, we must have
302 302 # killed it but not yet seen its return code. We don't wait for it,
303 303 # in case it's stuck in uninterruptible sleep. -9 = SIGKILL
304 304 rc = p.returncode or -9
305 305 raise CalledProcessError(rc, cell)
306 306
307 307 shebang.__skip_doctest__ = os.name != "posix"
308 308
309 309 async def _run_script(self, p, cell, to_close):
310 310 """callback for running the script in the background"""
311 311
312 312 p.stdin.write(cell)
313 313 await p.stdin.drain()
314 314 p.stdin.close()
315 315 await p.stdin.wait_closed()
316 316 await p.wait()
317 317 # asyncio read pipes have no close
318 318 # but we should drain the data anyway
319 319 for s in to_close:
320 320 await s.read()
321 321 self._gc_bg_processes()
322 322
323 323 @line_magic("killbgscripts")
324 324 def killbgscripts(self, _nouse_=''):
325 325 """Kill all BG processes started by %%script and its family."""
326 326 self.kill_bg_processes()
327 327 print("All background processes were killed.")
328 328
329 329 def kill_bg_processes(self):
330 330 """Kill all BG processes which are still running."""
331 331 if not self.bg_processes:
332 332 return
333 333 for p in self.bg_processes:
334 334 if p.returncode is None:
335 335 try:
336 336 p.send_signal(signal.SIGINT)
337 337 except:
338 338 pass
339 339 time.sleep(0.1)
340 340 self._gc_bg_processes()
341 341 if not self.bg_processes:
342 342 return
343 343 for p in self.bg_processes:
344 344 if p.returncode is None:
345 345 try:
346 346 p.terminate()
347 347 except:
348 348 pass
349 349 time.sleep(0.1)
350 350 self._gc_bg_processes()
351 351 if not self.bg_processes:
352 352 return
353 353 for p in self.bg_processes:
354 354 if p.returncode is None:
355 355 try:
356 356 p.kill()
357 357 except:
358 358 pass
359 359 self._gc_bg_processes()
360 360
361 361 def _gc_bg_processes(self):
362 362 self.bg_processes = [p for p in self.bg_processes if p.returncode is None]
General Comments 0
You need to be logged in to leave comments. Login now