demonstration of async communication between jupyter kernelgateway and aiohttp

gateway_websocket.py

#!/usr/bin/env python
"""
Simple demonstration of a aiohttp <-> Jupyter kernel gateway communication.

before running this script, the jupyter kernel gateway should be started with:

$ jupyter kernelgateway --JupyterWebsocketPersonality.list_kernels=True
"""

import aiohttp
import asyncio
import uuid
import json


def new_execute_msg(code):
    msg_id = uuid.uuid4().hex
    msg = {
        'header': {
            'username': '',
            'version': '5.0',
            'session': '',
            'msg_id': msg_id,
            'msg_type': 'execute_request'
        },
        'parent_header': {},
        'channel': 'shell',
        'content': {
            'code': code,
            'silent': False,
            'store_history': False,
            'user_expressions': {},
            'allow_stdin': False
        },
        'metadata': {},
        'buffers': {}
        }
    return msg


async def get_kernel_list(session, base_http_url):
    url = '{}/api/kernels'.format(base_http_url)
    print('get kernel list from: ', url)
    async with session.get(url) as resp:
        return await resp.json()


async def start_kernel(session, base_http_url, use_existing=False):
    url = '{}/api/kernels'.format(base_http_url)
    payload = {'name': 'python3'}

    kernel_list = await get_kernel_list(session, base_http_url)
    if use_existing and len(kernel_list) == 0:
        print('no kernel exists')
        use_existing = False

    if use_existing:
        kernel_id = kernel_list[0]['id']
        print('using existing kernel {}'.format(kernel_id))
    else:
        async with session.post(url, json=payload) as resp:
            body = await resp.json()
            kernel_id = body['id']
        print('starting new kernel {}'.format(kernel_id))
    return kernel_id


async def communicate(session, base_ws_url, kernel_id):
    channel_url = '{}/api/kernels/{}/channels'.format(base_ws_url, kernel_id)
    async with session.ws_connect(channel_url) as ws:
        msg = new_execute_msg('2+3')
        print('\n---- sending ----\n')
        print(msg)
        ws.send_json(msg)

        print('\n---- receiving ----\n')
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                msg_dict = json.loads(msg.data)
                msg_type = msg_dict['msg_type']
                content = msg_dict['content']
                print('< msg type: ' + msg_type + ' >')
                print(content)
                if msg_type == 'status' and \
                        content['execution_state'] == 'idle':
                    break


async def client_session():
    kernel_id = 'mykernel'
    auth = aiohttp.BasicAuth('fakeuser', 'fakepwd')
    base_http_url = 'http://localhost:8888'
    base_ws_url = 'ws://localhost:8888'
    async with aiohttp.ClientSession(auth=auth) as session:
        print('\n==== start or get kernel ====')
        kernel_id = await start_kernel(session, base_http_url,
                                       use_existing=True)
        print('\n==== start communication ====')
        await communicate(session, base_ws_url, kernel_id)


loop = asyncio.get_event_loop()
loop.run_until_complete(client_session())
loop.close()