File: jupyter.py

package info (click to toggle)
python-pweave 0.30.3-2
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 5,068 kB
  • sloc: python: 30,281; makefile: 167
file content (206 lines) | stat: -rw-r--r-- 6,780 bytes parent folder | download | duplicates (2)
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))