##// END OF EJS Templates
update execute preprocessor for msg spec 5
MinRK -
Show More
@@ -1,120 +1,162 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 #-----------------------------------------------------------------------------
7 # Imports
8 #-----------------------------------------------------------------------------
9
10 import os
6 import os
11 import sys
7 import sys
12
8
13 try:
9 try:
14 from queue import Empty # Py 3
10 from queue import Empty # Py 3
15 except ImportError:
11 except ImportError:
16 from Queue import Empty # Py 2
12 from Queue import Empty # Py 2
17
13
18 from IPython.kernel import KernelManager
14 from IPython.utils.traitlets import List, Unicode
19 from IPython.nbformat.current import reads, NotebookNode, writes
20
15
16 from IPython.nbformat.current import reads, NotebookNode, writes
21 from .base import Preprocessor
17 from .base import Preprocessor
22
18
19 # default timeout for reply and output: 10s
20 TIMEOUT = 10
23
21
24 #-----------------------------------------------------------------------------
25 # Classes
26 #-----------------------------------------------------------------------------
27 class ExecutePreprocessor(Preprocessor):
22 class ExecutePreprocessor(Preprocessor):
28 """
23 """
29 Executes all the cells in a notebook
24 Executes all the cells in a notebook
30 """
25 """
31 def __init__(self, extra_arguments=[], **kwargs):
26
32 """
27 # FIXME: to be removed with nbformat v4
33 Start an kernel to run the Python code
28 # map msg_type to v3 output_type
34 """
29 msg_type_map = {
35 super(ExecutePreprocessor, self).__init__(**kwargs)
30 "error" : "pyerr",
36 self.extra_arguments = []
31 "execute_result" : "pyout",
37
32 }
33
34 # FIXME: to be removed with nbformat v4
35 # map mime-type to v3 mime-type keys
36 mime_map = {
37 "text/plain" : "text",
38 "text/html" : "html",
39 "image/svg+xml" : "svg",
40 "image/png" : "png",
41 "image/jpeg" : "jpeg",
42 "text/latex" : "latex",
43 "application/json" : "json",
44 "application/javascript" : "javascript",
45 }
46
47 extra_arguments = List(Unicode)
48
38 def _create_client(self):
49 def _create_client(self):
50 from IPython.kernel import KernelManager
39 self.km = KernelManager()
51 self.km = KernelManager()
40 self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
52 self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
41 self.kc = self.km.client()
53 self.kc = self.km.client()
42 self.kc.start_channels()
54 self.kc.start_channels()
55 self.log.debug('kc.start_channels: %s', self.kc.session.session)
43 self.iopub = self.kc.iopub_channel
56 self.iopub = self.kc.iopub_channel
44 self.shell = self.kc.shell_channel
57 self.shell = self.kc.shell_channel
45 self.shell.kernel_info()
58 self.shell.kernel_info()
59 try:
60 self.shell.get_msg(timeout=TIMEOUT)
61 except Empty:
62 self.log.error("Timeout waiting for kernel_info reply")
63 raise
64 try:
65 self.iopub.get_msg(timeout=TIMEOUT)
66 except Empty:
67 self.log.warn("Timeout waiting for IOPub on startup")
46
68
47 def _shutdown_client(self):
69 def _shutdown_client(self):
48 self.kc.stop_channels()
70 self.kc.stop_channels()
49 self.km.shutdown_kernel()
71 self.km.shutdown_kernel()
50 del self.km
72 del self.km
51
73
52 def preprocess(self, nb, resources):
74 def preprocess(self, nb, resources):
53 self._create_client()
75 self._create_client()
54 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
76 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
55 self._shutdown_client()
77 self._shutdown_client()
56 return nb, resources
78 return nb, resources
57
79
58 def preprocess_cell(self, cell, resources, cell_index):
80 def preprocess_cell(self, cell, resources, cell_index):
59 """
81 """
60 Apply a transformation on each code cell. See base.py for details.
82 Apply a transformation on each code cell. See base.py for details.
61 """
83 """
62 if cell.cell_type != 'code':
84 if cell.cell_type != 'code':
63 return cell, resources
85 return cell, resources
64 try:
86 try:
65 outputs = self.run_cell(self.shell, self.iopub, cell)
87 outputs = self.run_cell(self.shell, self.iopub, cell)
66 except Exception as e:
88 except Exception as e:
67 self.log.error("failed to run cell: " + repr(e))
89 self.log.error("failed to run cell: " + repr(e))
68 self.log.error(str(cell.input))
90 self.log.error(str(cell.input))
69 sys.exit(1)
91 raise
70 cell.outputs = outputs
92 cell.outputs = outputs
71 return cell, resources
93 return cell, resources
72
94
73 @staticmethod
95 def run_cell(self, shell, iopub, cell):
74 def run_cell(shell, iopub, cell):
96 msg_id = shell.execute(cell.input)
75 # print cell.input
97 self.log.debug("Executing cell:\n%s", cell.input)
76 shell.execute(cell.input)
98 # wait for finish, with timeout
77 # wait for finish, maximum 20s
99 while True:
78 shell.get_msg(timeout=20)
100 try:
101 msg = shell.get_msg(timeout=TIMEOUT)
102 except Empty:
103 self.log.error("Timeout waiting for execute reply")
104 raise
105 if msg['parent_header'].get('msg_id') == msg_id:
106 break
107 else:
108 # not our reply
109 continue
110
79 outs = []
111 outs = []
80
112
81 while True:
113 while True:
82 try:
114 try:
83 msg = iopub.get_msg(timeout=0.2)
115 msg = iopub.get_msg(timeout=TIMEOUT)
84 except Empty:
116 except Empty:
117 self.log.warn("Timeout waiting for IOPub output")
85 break
118 break
119 if msg['parent_header'].get('msg_id') != msg_id:
120 # not an output from our execution
121 continue
86
122
87 msg_type = msg['msg_type']
123 msg_type = msg['msg_type']
124 self.log.debug("output: %s", msg_type)
88 content = msg['content']
125 content = msg['content']
89 out = NotebookNode(output_type=msg_type)
126 if msg_type == 'status':
127 if content['execution_state'] == 'idle':
128 break
129 else:
130 continue
131 elif msg_type in {'execute_input', 'pyin'}:
132 continue
133 elif msg_type == 'clear_output':
134 outs = []
135 continue
136
137 out = NotebookNode(output_type=self.msg_type_map.get(msg_type, msg_type))
90
138
91 # set the prompt number for the input and the output
139 # set the prompt number for the input and the output
92 if 'execution_count' in content:
140 if 'execution_count' in content:
93 cell['prompt_number'] = content['execution_count']
141 cell['prompt_number'] = content['execution_count']
94 out.prompt_number = content['execution_count']
142 out.prompt_number = content['execution_count']
95
143
96 if msg_type in ('status', 'pyin'):
97 continue
98 elif msg_type == 'clear_output':
99 outs = []
100 continue
101
102 if msg_type == 'stream':
144 if msg_type == 'stream':
103 out.stream = content['name']
145 out.stream = content['name']
104 out.text = content['data']
146 out.text = content['data']
105 elif msg_type in ('display_data', 'pyout'):
147 elif msg_type in ('display_data', 'execute_result'):
106 out['metadata'] = content['metadata']
148 out['metadata'] = content['metadata']
107 for mime, data in content['data'].iteritems():
149 for mime, data in content['data'].items():
108 attr = mime.split('/')[-1].lower()
150 # map mime-type keys to nbformat v3 keys
109 # this gets most right, but fix svg+html, plain
151 # this will be unnecessary in nbformat v4
110 attr = attr.replace('+xml', '').replace('plain', 'text')
152 key = self.mime_map.get(mime, mime)
111 setattr(out, attr, data)
153 out[key] = data
112 elif msg_type == 'pyerr':
154 elif msg_type == 'error':
113 out.ename = content['ename']
155 out.ename = content['ename']
114 out.evalue = content['evalue']
156 out.evalue = content['evalue']
115 out.traceback = content['traceback']
157 out.traceback = content['traceback']
116 else:
158 else:
117 self.log.error("unhandled iopub msg: " + msg_type)
159 self.log.error("unhandled iopub msg: " + msg_type)
118
160
119 outs.append(out)
161 outs.append(out)
120 return outs
162 return outs
@@ -1,87 +1,86 b''
1 """
1 """
2 Module with tests for the execute preprocessor.
2 Module with tests for the execute preprocessor.
3 """
3 """
4
4
5 # Copyright (c) IPython Development Team.
5 # Copyright (c) IPython Development Team.
6 # Distributed under the terms of the Modified BSD License.
6 # Distributed under the terms of the Modified BSD License.
7
7
8 import copy
8 import copy
9 import os
9 import os
10 import re
10 import re
11
11
12 from IPython.nbformat import current as nbformat
12 from IPython.nbformat import current as nbformat
13
13
14 from .base import PreprocessorTestsBase
14 from .base import PreprocessorTestsBase
15 from ..execute import ExecutePreprocessor
15 from ..execute import ExecutePreprocessor
16
16
17 from IPython.nbconvert.filters import strip_ansi
17 from IPython.nbconvert.filters import strip_ansi
18
18
19
19
20 addr_pat = re.compile(r'0x[0-9a-f]{7,9}')
20 addr_pat = re.compile(r'0x[0-9a-f]{7,9}')
21
21
22
23 class TestExecute(PreprocessorTestsBase):
22 class TestExecute(PreprocessorTestsBase):
24 """Contains test functions for execute.py"""
23 """Contains test functions for execute.py"""
25
24
26 @staticmethod
25 @staticmethod
27 def normalize_cell(cell):
26 def normalize_output(output):
28 """
27 """
29 Normalizes cells for comparison.
28 Normalizes outputs for comparison.
30 """
29 """
31 cell = dict(cell)
30 output = dict(output)
32 if 'metadata' in cell:
31 if 'metadata' in output:
33 del cell['metadata']
32 del output['metadata']
34 if 'text' in cell:
33 if 'text' in output:
35 cell['text'] = re.sub(addr_pat, '<HEXADDR>', cell['text'])
34 output['text'] = re.sub(addr_pat, '<HEXADDR>', output['text'])
36 if 'svg' in cell:
35 if 'svg' in output:
37 del cell['text']
36 del output['text']
38 if 'traceback' in cell:
37 if 'traceback' in output:
39 tb = []
38 tb = []
40 for line in cell['traceback']:
39 for line in output['traceback']:
41 tb.append(strip_ansi(line))
40 tb.append(strip_ansi(line))
42 cell['traceback'] = tb
41 output['traceback'] = tb
43
42
44 return cell
43 return output
45
44
46
45
47 def assert_notebooks_equal(self, expected, actual):
46 def assert_notebooks_equal(self, expected, actual):
48 expected_cells = expected['worksheets'][0]['cells']
47 expected_cells = expected['worksheets'][0]['cells']
49 actual_cells = actual['worksheets'][0]['cells']
48 actual_cells = actual['worksheets'][0]['cells']
50 assert len(expected_cells) == len(actual_cells)
49 assert len(expected_cells) == len(actual_cells)
51
50
52 for expected_cell, actual_cell in zip(expected_cells, actual_cells):
51 for expected_cell, actual_cell in zip(expected_cells, actual_cells):
53 expected_outputs = expected_cell.get('outputs', [])
52 expected_outputs = expected_cell.get('outputs', [])
54 actual_outputs = actual_cell.get('outputs', [])
53 actual_outputs = actual_cell.get('outputs', [])
55 normalized_expected_outputs = list(map(self.normalize_cell, expected_outputs))
54 normalized_expected_outputs = list(map(self.normalize_output, expected_outputs))
56 normalized_actual_outputs = list(map(self.normalize_cell, actual_outputs))
55 normalized_actual_outputs = list(map(self.normalize_output, actual_outputs))
57 assert normalized_expected_outputs == normalized_actual_outputs
56 assert normalized_expected_outputs == normalized_actual_outputs
58
57
59
58
60 def build_preprocessor(self):
59 def build_preprocessor(self):
61 """Make an instance of a preprocessor"""
60 """Make an instance of a preprocessor"""
62 preprocessor = ExecutePreprocessor()
61 preprocessor = ExecutePreprocessor()
63 preprocessor.enabled = True
62 preprocessor.enabled = True
64 return preprocessor
63 return preprocessor
65
64
66
65
67 def test_constructor(self):
66 def test_constructor(self):
68 """Can a ExecutePreprocessor be constructed?"""
67 """Can a ExecutePreprocessor be constructed?"""
69 self.build_preprocessor()
68 self.build_preprocessor()
70
69
71
70
72 def test_run_notebooks(self):
71 def test_run_notebooks(self):
73 """Runs a series of test notebooks and compares them to their actual output"""
72 """Runs a series of test notebooks and compares them to their actual output"""
74 current_dir = os.path.dirname(__file__)
73 current_dir = os.path.dirname(__file__)
75 input_files = os.listdir(os.path.join(current_dir, 'input'))
74 input_files = os.listdir(os.path.join(current_dir, 'input'))
76 for filename in input_files:
75 for filename in input_files:
77 if not filename.endswith(".ipynb"):
76 if not filename.endswith(".ipynb"):
78 continue
77 continue
79 with open(os.path.join(current_dir, 'input', filename)) as f:
78 with open(os.path.join(current_dir, 'input', filename)) as f:
80 input_nb = nbformat.read(f, 'ipynb')
79 input_nb = nbformat.read(f, 'ipynb')
81 with open(os.path.join(current_dir, 'expected', filename)) as f:
80 with open(os.path.join(current_dir, 'expected', filename)) as f:
82 expected_nb = nbformat.read(f, 'ipynb')
81 expected_nb = nbformat.read(f, 'ipynb')
83 res = self.build_resources()
82 res = self.build_resources()
84 preprocessor = self.build_preprocessor()
83 preprocessor = self.build_preprocessor()
85 output_nb, _ = preprocessor(input_nb, res)
84 output_nb, _ = preprocessor(input_nb, res)
86 self.assert_notebooks_equal(output_nb, expected_nb)
85 self.assert_notebooks_equal(output_nb, expected_nb)
87
86
General Comments 0
You need to be logged in to leave comments. Login now