-
Notifications
You must be signed in to change notification settings - Fork 64
/
jupyter.py
206 lines (166 loc) · 6.62 KB
/
jupyter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
# -*- coding: utf-8 -*-
from jupyter_client.manager import start_new_kernel
from jupyter_client import KernelManager
from nbformat.v4 import output_from_msg
import os
from .. import config
from .base import PwebProcessorBase
from . import subsnippets
from IPython.core import inputsplitter
from ipykernel.inprocess import InProcessKernelManager
from queue import Empty
class JupyterProcessor(PwebProcessorBase):
"""Generic Jupyter processor, should work with any kernel"""
def __init__(self, parsed, kernel, source, mode,
figdir, outdir, embed_kernel=False):
super(JupyterProcessor, self).__init__(parsed, kernel, source, mode,
figdir, outdir)
self.extra_arguments = None
self.timeout = -1
path = os.path.abspath(outdir)
if embed_kernel:
km = InProcessKernelManager(kernel_name=kernel)
else:
km = KernelManager(kernel_name=kernel)
km.start_kernel(cwd=path, stderr=open(os.devnull, 'w'))
kc = km.client()
kc.start_channels()
try:
kc.wait_for_ready()
except RuntimeError:
print("Timeout from starting kernel\nTry restarting python session and running weave again")
kc.stop_channels()
km.shutdown_kernel()
raise
self.km = km
self.kc = kc
self.kc.allow_stdin = False
def close(self):
self.kc.stop_channels()
self.km.shutdown_kernel()
def run_cell(self, src):
cell = {}
cell["source"] = src.lstrip()
msg_id = self.kc.execute(src.lstrip(), store_history=False)
# wait for finish, with timeout
while True:
try:
timeout = self.timeout
if timeout < 0:
timeout = None
msg = self.kc.get_shell_msg(timeout=timeout)
except Empty:
if self.interrupt_on_timeout:
self.km.interrupt_kernel()
break
else:
try:
exception = TimeoutError
except NameError:
exception = RuntimeError
raise exception(
"Cell execution timed out, see log for details.")
if msg['parent_header'].get('msg_id') == msg_id:
break
else:
# not our reply
continue
outs = []
while True:
try:
# We've already waited for execute_reply, so all output
# should already be waiting. However, on slow networks, like
# in certain CI systems, waiting < 1 second might miss messages.
# So long as the kernel sends a status:idle message when it
# finishes, we won't actually have to wait this long, anyway.
msg = self.kc.iopub_channel.get_msg(block=True, timeout=4)
except Empty:
print("Timeout waiting for IOPub output\nTry restarting python session and running weave again")
raise RuntimeError("Timeout waiting for IOPub output")
#stdout from InProcessKernelManager has no parent_header
if msg['parent_header'].get('msg_id') != msg_id and msg['msg_type'] != "stream":
continue
msg_type = msg['msg_type']
content = msg['content']
# set the prompt number for the input and the output
if 'execution_count' in content:
cell['execution_count'] = content['execution_count']
if msg_type == 'status':
if content['execution_state'] == 'idle':
break
else:
continue
elif msg_type == 'execute_input':
continue
elif msg_type == 'clear_output':
outs = []
continue
elif msg_type.startswith('comm'):
continue
try:
out = output_from_msg(msg)
except ValueError:
print("unhandled iopub msg: " + msg_type)
else:
outs.append(out)
return outs
def loadstring(self, code_str, **kwargs):
return self.run_cell(code_str)
#Yes same format for compatibility even if term is not implemented
def loadterm(self, code_str, **kwargs):
return((sources, self.run_cell(code_str)))
#TODO add support for "rich" output
#Requires storing the results for formatter
def load_inline_string(self, code_string):
from nbconvert import filters
outputs = self.loadstring(code_string)
result = ""
for out in outputs:
if out["output_type"] == "stream":
result += out["text"]
elif out["output_type"] == "error":
result += filters.strip_ansi("".join(out["traceback"]))
elif "text/plain" in out["data"]:
result += out["data"]["text/plain"]
else:
result = ""
return result
class IPythonProcessor(JupyterProcessor):
"""Contains IPython specific functions"""
def __init__(self, *args):
kernel = args[1]
if kernel == "python3":
embed = True
else:
embed = False
super(IPythonProcessor, self).__init__(*args, embed_kernel=embed)
if config.rcParams["usematplotlib"]:
self.init_matplotlib()
def init_matplotlib(self):
self.loadstring(subsnippets.init_matplotlib)
def pre_run_hook(self, chunk):
f_size = """matplotlib.rcParams.update({"figure.figsize" : (%i, %i)})""" % chunk["f_size"]
f_dpi = """matplotlib.rcParams.update({"figure.dpi" : %i})""" % chunk["dpi"]
self.loadstring("\n".join([f_size, f_dpi]))
def loadterm(self, code_str, **kwargs):
splitter = inputsplitter.IPythonInputSplitter()
code_lines = code_str.lstrip().splitlines()
sources = []
outputs = []
for line in code_lines:
if splitter.push_accepts_more():
splitter.push(line)
else:
code_str = splitter.source
sources.append(code_str)
out = self.loadstring(code_str)
#print(out)
outputs.append(out)
splitter.reset()
splitter.push(line)
if splitter.source != "":
code_str = splitter.source
sources.append(code_str)
out = self.loadstring(code_str)
outputs.append(out)
return((sources, outputs))