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