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