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