,python

#!/bin/python

import os
import signal
import subprocess
import sys
import select
import jupyter_client
from time import sleep

usage = """# Usage:

Execute command (and start a kernel if one isn't running):
    ,python [command]

Kill running kernel:
    ,python --kill

New kernel (and kill any existing):
    ,python --new

New kernel (and kill any existing) and execute command:
    ,python --new [command]

Jump into ipython shell of existing kernel:
    ,python --shell

"""

to_execute = None
kill = False
restart = False

try:
    arg = sys.argv[1]
    if arg in ['-h', '--help']:
        print(usage)
        exit(0)
    elif arg == '--kill':
        kill = True
    elif arg == '--new':
        restart = True
        try:
            to_execute = sys.argv[2]
        except:
            pass
    elif arg in ['--shell', '--console']:
        os.system('jupyter console --existing shell_kernel.json')
        sys.exit(0)
    elif arg in ['--qtconsole']:
        os.system('jupyter qtconsole --existing shell_kernel.json')
        sys.exit(0)
    else:
        to_execute = arg
except SystemExit:
    pass
except:
    print(usage)
    exit(1)

connection_file = 'shell_kernel.json'
pid_file = 'shell_kernel.pid'
runtime_path = os.path.expanduser('~/.local/share/jupyter/runtime/')
connection_file_path = os.path.join(runtime_path, connection_file)
connection_pid_path =  os.path.join(runtime_path, pid_file)

# If the pid file exists, check if it exists and is a jupyter-kernel as expected
if os.path.exists(connection_pid_path):
    with open(connection_pid_path) as f:
        pid = int(f.read().strip())
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        kernel_exists = False
    else:
        # pid exists, but confirm that it is the kernel and not something else
        # e.g. case of rebooted, pid file still exists, different process now
        cmd = subprocess.check_output(["ps", "-w", "-w", "p", str(pid), "-o", "cmd", "h"]).decode()
        kernel_exists = connection_file_path in cmd and 'jupyter-kernel' in cmd
else:
    kernel_exists = False    

# Kill or restart the kernel if requested, if we confirmed the pid is a jupyter-kernel
if kill or restart:
    if kernel_exists:
        os.kill(pid, signal.SIGTERM)
        print(f"Killed kernel with pid {pid}", file=sys.stderr)
        kernel_exists = False
    else:
        print("No kernel found to kill", file=sys.stderr)

# Cleanup any dead connection files
if not kernel_exists and os.path.exists(connection_file_path):
    os.remove(connection_file_path)
if not kernel_exists and os.path.exists(connection_pid_path):
    os.remove(connection_pid_path)

if kill:
    sys.exit(0)

# Creating a new kernel
is_new_kernel = False
if not kernel_exists:
    cf_arg = '--KernelManager.connection_file=' + connection_file_path
    proc = subprocess.Popen(['jupyter', 'kernel', cf_arg], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True)
    pid = str(proc.pid)
    print(f"Started kernel with pid {pid}", file=sys.stderr)
    with open(connection_pid_path, 'w') as f:
        f.write(pid)
    is_new_kernel = True
    # HACK to avoid km.iopub_channel.msg_ready() indefinitely False
    # sleep(1)

# Find the connection file, waiting max 2 seconds before exit(1)
cf = None
max_try = 10
i = 0
while True:
    try:
        cf = jupyter_client.find_connection_file(connection_file_path)
        break
    except:
        i += 1
        sleep(0.2)
        if i >= max_try:
            print("Couldn't find connection file")
            sys.exit(1)


# Create our client
# Inspiration from https://stackoverflow.com/a/35792945
km = jupyter_client.BlockingKernelClient(connection_file=cf)
km.load_connection_file()


def get_messages(km, print_output=True):
    busy_seen = False  # Ensure we see a "busy" execution_state before we look for "idle"
    wait_for_idle = True  # Stop reading messages once we see "idle" (after "busy")
    wait_count = 0
    max_wait = 20  # Don't spin indefinitely if kernel is not giving us any messages
    sleep_time = 0.1
    msgs = []
    while wait_for_idle:
        if km.iopub_channel.msg_ready():
            msg = km.iopub_channel.get_msg()
            msgs.append(msg)
            if 'content' not in msg:
                continue
            content = msg['content']
            if 'execution_state' in content:
                if content['execution_state'] == 'busy':
                    busy_seen = True
                if content['execution_state'] == 'idle' and busy_seen:
                    wait_for_idle = False
            if 'name' in content and 'text' in content:
                name = content['name']
                text = content['text']
                if name == 'stdout':
                    print(text, end='')
                elif name == 'stderr':
                    print(text, end='', file=sys.stderr)
            if 'data' in content:
                for d in content['data'].values():
                    print(d)
            if 'traceback' in content:
                for t in content['traceback']:
                    print(t, file=sys.stderr)
        elif not len(msgs):
            wait_count += 1
            if wait_count > max_wait:
                print("No messages from kernel")
                return []
        sleep(sleep_time)
    return msgs


def exec_km(to_execute):
    # TODO is km.execute_interactive more suitable?
    km.execute(to_execute)
    return get_messages(km)


def stdin_avail():
    return select.select([sys.stdin, ], [], [], 0.0)[0]


# Request kernel info before executing anything
km.kernel_info()
if not get_messages(km):
    print("Couldn't get kernel info")
    exit(1)

# On first run of a new kernel, you may want to execute certain lines
# e.g. `import os`, `import my_library`, or even `%pylab`
if is_new_kernel:
    # exec_km('%pylab')
    exec_km("""
import os
import re
_quit = quit
def quit():
    _quit(keep_kernel=True)
exit = quit
""")
    # _quit = quit
    # def quit():
    #     _quit(keep_kernel=True)
    # print(dir(km))
    # km.shell.push({'_quit' : _quit, 'quit' : quit, 'exit' : quit})

if to_execute is None:
    exit(0)

# Assign any stdin to `stdin` and `lines` variables
if stdin_avail():
    from_stdin = sys.stdin.read()
    exec_km('stdin = ' + from_stdin.__repr__())
    exec_km('lines = stdin.splitlines()')

# Execute the requested command from the terminal
exec_km(to_execute)