tst.py

"""Running code with jupyter_client and saving outputs"""

# LICENSE: CC-0 (public domain)

from binascii import a2b_base64

from jupyter_client.manager import start_new_kernel

def handle_output(msg):
    msg_type = msg['header']['msg_type']
    if msg_type in {'display_data', 'execute_result'}:
        data = msg['content']['data']
        print("%s: %s" % (msg_type, data.get('text/plain', 'unknown')))
        print("mimetypes: %s" % sorted(data.keys()))

km, kc = start_new_kernel()

# capture outputs and show info about them
kc.execute_interactive("""
%matplotlib inline
import matplotlib.pyplot as plt
plt.plot([1,3,2])
""", output_hook=handle_output)

# save PNG outputs to files
saved_images = []
tpl = 'plot_%i.png'
def save_images(msg):
    msg_type = msg['header']['msg_type']
    if msg_type in {'display_data', 'execute_result'}:
        data = msg['content']['data']
        b64_png = data.get('image/png')
        if not b64_png:
            return
        png_bytes = a2b_base64(b64_png)
        fname = tpl % len(saved_images)
        saved_images.append(fname)
        print("Saving %s to %s" % (data.get('text/plain', 'unknown'), fname))
        with open(fname, 'wb') as f:
            f.write(png_bytes)

kc.execute_interactive("""
import numpy as np
x = np.linspace(0, 10)
plt.plot(x, np.sin(x))
plt.figure()
plt.plot(x, np.cos(x))
""", output_hook=save_images)

print("Saved", ', '.join(saved_images))

km.shutdown_kernel()
km.cleanup()