#! /usr/bin/env python
"""
A script for watching all traffic on the IOPub channel (stdout/stderr/pyerr) of
engines.

This connects to the default cluster, or you can pass the path to your
ipcontroller-client.json

Try running this script, and then running a few jobs that print (and call
sys.stdout.flush), and you will see the print statements as they arrive,
notably not waiting for the results to finish.

You can use the zeromq SUBSCRIBE mechanism to only receive information from
specific engines, and easily filter by message type.

Authors
-------
* MinRK
"""

import sys
import json
import zmq

from IPython.kernel.zmq.session import Session
from IPython.utils.py3compat import str_to_bytes
from IPython.utils.path import get_security_file


def main(connection_file):
    """watch iopub channel, and print messages"""

    ctx = zmq.Context.instance()

    with open(connection_file) as f:
        cfg = json.loads(f.read())

    reg_url = cfg['interface']
    iopub_port = cfg['iopub']
    iopub_url = "%s:%s" % (reg_url, iopub_port)

    session = Session(key=str_to_bytes(cfg['key']))
    sub = ctx.socket(zmq.SUB)

    # This will subscribe to all messages:
    sub.setsockopt(zmq.SUBSCRIBE, b'')
    # replace with b'' with b'engine.1.stdout' to subscribe only to engine 1's
    # stdout 0MQ subscriptions are simple 'foo*' matches, so 'engine.1.'
    # subscribes to everything from engine 1, but there is no way to subscribe
    # to just stdout from everyone.  multiple calls to subscribe will add
    # subscriptions, e.g. to subscribe to engine 1's stderr and engine 2's
    # stdout: sub.setsockopt(zmq.SUBSCRIBE, b'engine.1.stderr')
    # sub.setsockopt(zmq.SUBSCRIBE, b'engine.2.stdout')
    sub.connect(iopub_url)
    while True:
        try:
            idents, msg = session.recv(sub, mode=0)
        except KeyboardInterrupt:
            return
        # ident always length 1 here
        topic = idents[0]
        if msg['msg_type'] == 'stream':
            # stdout/stderr
            # stream names are in msg['content']['name'], if you want to handle
            # them differently
            print("%s: %s" % (topic, msg['content']['data']))
        elif msg['msg_type'] == 'pyerr':
            # Python traceback
            c = msg['content']
            print(topic + ':')
            for line in c['traceback']:
                # indent lines
                print('    ' + line)

if __name__ == '__main__':
    if len(sys.argv) > 1:
        cf = sys.argv[1]
    else:
        # This gets the security file for the default profile:
        cf = get_security_file('ipcontroller-client.json')
    main(cf)

Add iopubwatcher.py to utils/