##// END OF EJS Templates
Fix queue import for Python 3
Julia Evans -
Show More
@@ -1,114 +1,118 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 #-----------------------------------------------------------------------------
6 #-----------------------------------------------------------------------------
7 # Imports
7 # Imports
8 #-----------------------------------------------------------------------------
8 #-----------------------------------------------------------------------------
9
9
10 import os
10 import os
11 import sys
11 import sys
12
12
13 from Queue import Empty
13 try:
14 from queue import Empty # Py 3
15 except ImportError:
16 from Queue import Empty # Py 2
17
14 from IPython.kernel import KernelManager
18 from IPython.kernel import KernelManager
15 from IPython.nbformat.current import reads, NotebookNode, writes
19 from IPython.nbformat.current import reads, NotebookNode, writes
16
20
17 from .base import Preprocessor
21 from .base import Preprocessor
18
22
19
23
20 #-----------------------------------------------------------------------------
24 #-----------------------------------------------------------------------------
21 # Classes
25 # Classes
22 #-----------------------------------------------------------------------------
26 #-----------------------------------------------------------------------------
23 class ExecutePreprocessor(Preprocessor):
27 class ExecutePreprocessor(Preprocessor):
24 """
28 """
25 Executes all the cells in a notebook
29 Executes all the cells in a notebook
26 """
30 """
27 def __init__(self, extra_arguments=[], **kwargs):
31 def __init__(self, extra_arguments=[], **kwargs):
28 """
32 """
29 Start an kernel to run the Python code
33 Start an kernel to run the Python code
30 """
34 """
31 super(ExecutePreprocessor, self).__init__(**kwargs)
35 super(ExecutePreprocessor, self).__init__(**kwargs)
32 self.extra_arguments = []
36 self.extra_arguments = []
33
37
34 def _create_client(self):
38 def _create_client(self):
35 self.km = KernelManager()
39 self.km = KernelManager()
36 self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
40 self.km.start_kernel(extra_arguments=self.extra_arguments, stderr=open(os.devnull, 'w'))
37 self.kc = self.km.client()
41 self.kc = self.km.client()
38 self.kc.start_channels()
42 self.kc.start_channels()
39 self.iopub = self.kc.iopub_channel
43 self.iopub = self.kc.iopub_channel
40 self.shell = self.kc.shell_channel
44 self.shell = self.kc.shell_channel
41 self.shell.kernel_info()
45 self.shell.kernel_info()
42
46
43 def _shutdown_client(self):
47 def _shutdown_client(self):
44 self.kc.stop_channels()
48 self.kc.stop_channels()
45 self.km.shutdown_kernel()
49 self.km.shutdown_kernel()
46 del self.km
50 del self.km
47
51
48 def preprocess(self, nb, resources):
52 def preprocess(self, nb, resources):
49 self._create_client()
53 self._create_client()
50 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
54 nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)
51 self._shutdown_client()
55 self._shutdown_client()
52 return nb, resources
56 return nb, resources
53
57
54 def preprocess_cell(self, cell, resources, cell_index):
58 def preprocess_cell(self, cell, resources, cell_index):
55 """
59 """
56 Apply a transformation on each code cell. See base.py for details.
60 Apply a transformation on each code cell. See base.py for details.
57 """
61 """
58 if cell.cell_type != 'code':
62 if cell.cell_type != 'code':
59 return cell, resources
63 return cell, resources
60 try:
64 try:
61 outputs = self.run_cell(self.shell, self.iopub, cell)
65 outputs = self.run_cell(self.shell, self.iopub, cell)
62 except Exception as e:
66 except Exception as e:
63 self.log.error("failed to run cell: " + repr(e))
67 self.log.error("failed to run cell: " + repr(e))
64 self.log.error(str(cell.input))
68 self.log.error(str(cell.input))
65 sys.exit(1)
69 sys.exit(1)
66 cell.outputs = outputs
70 cell.outputs = outputs
67 return cell, resources
71 return cell, resources
68
72
69 @staticmethod
73 @staticmethod
70 def run_cell(shell, iopub, cell):
74 def run_cell(shell, iopub, cell):
71 # print cell.input
75 # print cell.input
72 shell.execute(cell.input)
76 shell.execute(cell.input)
73 # wait for finish, maximum 20s
77 # wait for finish, maximum 20s
74 shell.get_msg(timeout=20)
78 shell.get_msg(timeout=20)
75 outs = []
79 outs = []
76
80
77 while True:
81 while True:
78 try:
82 try:
79 msg = iopub.get_msg(timeout=0.2)
83 msg = iopub.get_msg(timeout=0.2)
80 except Empty:
84 except Empty:
81 break
85 break
82 msg_type = msg['msg_type']
86 msg_type = msg['msg_type']
83 if msg_type in ('status', 'pyin'):
87 if msg_type in ('status', 'pyin'):
84 continue
88 continue
85 elif msg_type == 'clear_output':
89 elif msg_type == 'clear_output':
86 outs = []
90 outs = []
87 continue
91 continue
88
92
89 content = msg['content']
93 content = msg['content']
90 # print msg_type, content
94 # print msg_type, content
91 out = NotebookNode(output_type=msg_type)
95 out = NotebookNode(output_type=msg_type)
92
96
93 if msg_type == 'stream':
97 if msg_type == 'stream':
94 out.stream = content['name']
98 out.stream = content['name']
95 out.text = content['data']
99 out.text = content['data']
96 elif msg_type in ('display_data', 'pyout'):
100 elif msg_type in ('display_data', 'pyout'):
97 out['metadata'] = content['metadata']
101 out['metadata'] = content['metadata']
98 for mime, data in content['data'].iteritems():
102 for mime, data in content['data'].iteritems():
99 attr = mime.split('/')[-1].lower()
103 attr = mime.split('/')[-1].lower()
100 # this gets most right, but fix svg+html, plain
104 # this gets most right, but fix svg+html, plain
101 attr = attr.replace('+xml', '').replace('plain', 'text')
105 attr = attr.replace('+xml', '').replace('plain', 'text')
102 setattr(out, attr, data)
106 setattr(out, attr, data)
103 if msg_type == 'pyout':
107 if msg_type == 'pyout':
104 out.prompt_number = content['execution_count']
108 out.prompt_number = content['execution_count']
105 elif msg_type == 'pyerr':
109 elif msg_type == 'pyerr':
106 out.ename = content['ename']
110 out.ename = content['ename']
107 out.evalue = content['evalue']
111 out.evalue = content['evalue']
108 out.traceback = content['traceback']
112 out.traceback = content['traceback']
109 else:
113 else:
110 self.log.error("unhandled iopub msg: " + msg_type)
114 self.log.error("unhandled iopub msg: " + msg_type)
111
115
112 outs.append(out)
116 outs.append(out)
113 return outs
117 return outs
114
118
General Comments 0
You need to be logged in to leave comments. Login now