##// END OF EJS Templates
Simplify ExecutePreprocessor using new context manager...
Thomas Kluyver -
Show More
@@ -1,165 +1,142 b''
1 """Module containing a preprocessor that removes the outputs from code cells"""
1 """Module containing a preprocessor that removes the outputs from code cells"""
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 os
6 import os
7 import sys
7 import sys
8
8
9 try:
9 try:
10 from queue import Empty # Py 3
10 from queue import Empty # Py 3
11 except ImportError:
11 except ImportError:
12 from Queue import Empty # Py 2
12 from Queue import Empty # Py 2
13
13
14 from IPython.utils.traitlets import List, Unicode
14 from IPython.utils.traitlets import List, Unicode
15
15
16 from IPython.nbformat.current import reads, NotebookNode, writes
16 from IPython.nbformat.current import reads, NotebookNode, writes
17 from .base import Preprocessor
17 from .base import Preprocessor
18 from IPython.utils.traitlets import Integer
18 from IPython.utils.traitlets import Integer
19
19
20 class ExecutePreprocessor(Preprocessor):
20 class ExecutePreprocessor(Preprocessor):
21 """
21 """
22 Executes all the cells in a notebook
22 Executes all the cells in a notebook
23 """
23 """
24
24
25 timeout = Integer(30, config=True,
25 timeout = Integer(30, config=True,
26 help="The time to wait (in seconds) for output from executions."
26 help="The time to wait (in seconds) for output from executions."
27 )
27 )
28 # FIXME: to be removed with nbformat v4
28 # FIXME: to be removed with nbformat v4
29 # map msg_type to v3 output_type
29 # map msg_type to v3 output_type
30 msg_type_map = {
30 msg_type_map = {
31 "error" : "pyerr",
31 "error" : "pyerr",
32 "execute_result" : "pyout",
32 "execute_result" : "pyout",
33 }
33 }
34
34
35 # FIXME: to be removed with nbformat v4
35 # FIXME: to be removed with nbformat v4
36 # map mime-type to v3 mime-type keys
36 # map mime-type to v3 mime-type keys
37 mime_map = {
37 mime_map = {
38 "text/plain" : "text",
38 "text/plain" : "text",
39 "text/html" : "html",
39 "text/html" : "html",
40 "image/svg+xml" : "svg",
40 "image/svg+xml" : "svg",
41 "image/png" : "png",
41 "image/png" : "png",
42 "image/jpeg" : "jpeg",
42 "image/jpeg" : "jpeg",
43 "text/latex" : "latex",
43 "text/latex" : "latex",
44 "application/json" : "json",
44 "application/json" : "json",
45 "application/javascript" : "javascript",
45 "application/javascript" : "javascript",
46 }
46 }
47
47
48 extra_arguments = List(Unicode)
48 extra_arguments = List(Unicode)
49
49
50 def _create_client(self):
51 from IPython.kernel import KernelManager
52 self.km = KernelManager()
53 self.km.write_connection_file()
54 self.kc = self.km.client()
55 self.kc.start_channels()
56 self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
57 self.iopub = self.kc.iopub_channel
58 self.shell = self.kc.shell_channel
59 self.shell.kernel_info()
60 try:
61 self.shell.get_msg(timeout=self.timeout)
62 except Empty:
63 self.log.error("Timeout waiting for kernel_info reply")
64 raise
65 # flush IOPub
66 while True:
67 try:
68 self.iopub.get_msg(block=True, timeout=0.25)
69 except Empty:
70 break
71
72 def _shutdown_client(self):
73 self.kc.stop_channels()
74 self.km.shutdown_kernel()
75 del self.km
76
77 def preprocess(self, nb, resources):
50 def preprocess(self, nb, resources):
78 self._create_client()
51 from IPython.kernel import run_kernel
52 kernel_name = nb.metadata.get('kernelspec', {}).get('name', 'python')
53 with run_kernel(kernel_name=kernel_name,
54 extra_arguments=self.extra_arguments,
55 stderr=open(os.devnull, 'w')) as kc:
56 self.kc = kc
79 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
57 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
80 self._shutdown_client()
81 return nb, resources
58 return nb, resources
82
59
83 def preprocess_cell(self, cell, resources, cell_index):
60 def preprocess_cell(self, cell, resources, cell_index):
84 """
61 """
85 Apply a transformation on each code cell. See base.py for details.
62 Apply a transformation on each code cell. See base.py for details.
86 """
63 """
87 if cell.cell_type != 'code':
64 if cell.cell_type != 'code':
88 return cell, resources
65 return cell, resources
89 try:
66 try:
90 outputs = self.run_cell(self.shell, self.iopub, cell)
67 outputs = self.run_cell(self.kc.shell_channel, self.kc.iopub_channel, cell)
91 except Exception as e:
68 except Exception as e:
92 self.log.error("failed to run cell: " + repr(e))
69 self.log.error("failed to run cell: " + repr(e))
93 self.log.error(str(cell.input))
70 self.log.error(str(cell.input))
94 raise
71 raise
95 cell.outputs = outputs
72 cell.outputs = outputs
96 return cell, resources
73 return cell, resources
97
74
98 def run_cell(self, shell, iopub, cell):
75 def run_cell(self, shell, iopub, cell):
99 msg_id = shell.execute(cell.input)
76 msg_id = shell.execute(cell.input)
100 self.log.debug("Executing cell:\n%s", cell.input)
77 self.log.debug("Executing cell:\n%s", cell.input)
101 # wait for finish, with timeout
78 # wait for finish, with timeout
102 while True:
79 while True:
103 try:
80 try:
104 msg = shell.get_msg(timeout=self.timeout)
81 msg = shell.get_msg(timeout=self.timeout)
105 except Empty:
82 except Empty:
106 self.log.error("Timeout waiting for execute reply")
83 self.log.error("Timeout waiting for execute reply")
107 raise
84 raise
108 if msg['parent_header'].get('msg_id') == msg_id:
85 if msg['parent_header'].get('msg_id') == msg_id:
109 break
86 break
110 else:
87 else:
111 # not our reply
88 # not our reply
112 continue
89 continue
113
90
114 outs = []
91 outs = []
115
92
116 while True:
93 while True:
117 try:
94 try:
118 msg = iopub.get_msg(timeout=self.timeout)
95 msg = iopub.get_msg(timeout=self.timeout)
119 except Empty:
96 except Empty:
120 self.log.warn("Timeout waiting for IOPub output")
97 self.log.warn("Timeout waiting for IOPub output")
121 break
98 break
122 if msg['parent_header'].get('msg_id') != msg_id:
99 if msg['parent_header'].get('msg_id') != msg_id:
123 # not an output from our execution
100 # not an output from our execution
124 continue
101 continue
125
102
126 msg_type = msg['msg_type']
103 msg_type = msg['msg_type']
127 self.log.debug("output: %s", msg_type)
104 self.log.debug("output: %s", msg_type)
128 content = msg['content']
105 content = msg['content']
129 if msg_type == 'status':
106 if msg_type == 'status':
130 if content['execution_state'] == 'idle':
107 if content['execution_state'] == 'idle':
131 break
108 break
132 else:
109 else:
133 continue
110 continue
134 elif msg_type in {'execute_input', 'pyin'}:
111 elif msg_type in {'execute_input', 'pyin'}:
135 continue
112 continue
136 elif msg_type == 'clear_output':
113 elif msg_type == 'clear_output':
137 outs = []
114 outs = []
138 continue
115 continue
139
116
140 out = NotebookNode(output_type=self.msg_type_map.get(msg_type, msg_type))
117 out = NotebookNode(output_type=self.msg_type_map.get(msg_type, msg_type))
141
118
142 # set the prompt number for the input and the output
119 # set the prompt number for the input and the output
143 if 'execution_count' in content:
120 if 'execution_count' in content:
144 cell['prompt_number'] = content['execution_count']
121 cell['prompt_number'] = content['execution_count']
145 out.prompt_number = content['execution_count']
122 out.prompt_number = content['execution_count']
146
123
147 if msg_type == 'stream':
124 if msg_type == 'stream':
148 out.stream = content['name']
125 out.stream = content['name']
149 out.text = content['data']
126 out.text = content['data']
150 elif msg_type in ('display_data', 'execute_result'):
127 elif msg_type in ('display_data', 'execute_result'):
151 out['metadata'] = content['metadata']
128 out['metadata'] = content['metadata']
152 for mime, data in content['data'].items():
129 for mime, data in content['data'].items():
153 # map mime-type keys to nbformat v3 keys
130 # map mime-type keys to nbformat v3 keys
154 # this will be unnecessary in nbformat v4
131 # this will be unnecessary in nbformat v4
155 key = self.mime_map.get(mime, mime)
132 key = self.mime_map.get(mime, mime)
156 out[key] = data
133 out[key] = data
157 elif msg_type == 'error':
134 elif msg_type == 'error':
158 out.ename = content['ename']
135 out.ename = content['ename']
159 out.evalue = content['evalue']
136 out.evalue = content['evalue']
160 out.traceback = content['traceback']
137 out.traceback = content['traceback']
161 else:
138 else:
162 self.log.error("unhandled iopub msg: " + msg_type)
139 self.log.error("unhandled iopub msg: " + msg_type)
163
140
164 outs.append(out)
141 outs.append(out)
165 return outs
142 return outs
General Comments 0
You need to be logged in to leave comments. Login now