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