IPython/kernel/zmq/iostream.py

import sys
import time
import os
from io import StringIO

from session import extract_header, Message

from IPython.utils import io, text
from IPython.utils import py3compat

import multiprocessing as mp
import multiprocessing.sharedctypes as mpshc
from ctypes import c_bool
#-----------------------------------------------------------------------------
# Globals
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Stream classes
#-----------------------------------------------------------------------------

class OutStream(object):
    """A file like object that publishes the stream to a 0MQ PUB socket."""

    # The time interval between automatic flushes, in seconds.
    flush_interval = 0.05
    topic=None

    def __init__(self, session, pub_socket, name):
        self.encoding = 'UTF-8'
        self.session = session
        self.pub_socket = pub_socket
        self.name = name
        self.parent_header = {}
        self._new_buffer()
        self._manager = mp.Manager()
        #use sharectype here so it don't have to hit the manager
        #no synchronize needed either(right?). Just a flag telling the master
        #to switch the buffer to que
        self._found_newprocess = mpshc.RawValue(c_bool, False)
        self._que_buffer = self._manager.Queue()
        self._que_lock = self._manager.Lock()
        self._masterpid = os.getpid()
        self._master_has_switched = False

    def _switch_to_que(self):
        #should only be called on master process
        #don't clear the que before putting data in since
        #child process might have put something in the que before the
        #master know it.
        self._que_buffer.put(self._buffer.getvalue())
        self._new_buffer()
        self._start = -1

    def _is_master_process(self):
        return os.getpid()==self._masterpid

    def _debug_print(self,s):
        sys.__stdout__.write(s+'\n')
        sys.__stdout__.flush()

    def _check_mp_mode(self):
        """check multiprocess and switch to que if necessary"""
        if not self._found_newprocess.value:
            if not self._is_master_process():
                self._found_newprocess.value = True
        elif self._found_newprocess.value and not self._master_has_switched:

            #switch to que if it has not been switch
            if self._is_master_process():
                self._switch_to_que()
                self._master_has_switched = True

        return self._found_newprocess.value


    def set_parent(self, parent):
        self.parent_header = extract_header(parent)

    def close(self):
        self.pub_socket = None

    def flush(self):
        #io.rprint('>>>flushing output buffer: %s<<<' % self.name)  # dbg

        if self.pub_socket is None:
            raise ValueError(u'I/O operation on closed file')
        else:
            if self._is_master_process():
                data = u''
                #obtain data
                if self._check_mp_mode():#multiprocess
                    with self._que_lock:
                        while not self._que_buffer.empty():
                            data += self._que_buffer.get()
                else:#single process mode
                    data = self._buffer.getvalue()

                if data:
                    content = {u'name':self.name, u'data':data}
                    msg = self.session.send(self.pub_socket, u'stream', content=content,
                                           parent=self.parent_header, ident=self.topic)

                    if hasattr(self.pub_socket, 'flush'):
                        # socket itself has flush (presumably ZMQStream)
                        self.pub_socket.flush()
                    self._buffer.close()
                    self._new_buffer()
            else:
                pass


    def isatty(self):
        return False

    def __next__(self):
        raise IOError('Read not supported on a write only stream.')

    if not py3compat.PY3:
        next = __next__

    def read(self, size=-1):
        raise IOError('Read not supported on a write only stream.')

    def readline(self, size=-1):
        raise IOError('Read not supported on a write only stream.')

    def write(self, string):
        if self.pub_socket is None:
            raise ValueError('I/O operation on closed file')
        else:
            # Make sure that we're handling unicode
            if not isinstance(string, unicode):
                string = string.decode(self.encoding, 'replace')

            if self._check_mp_mode(): #multi process mode
                with self._que_lock:
                    self._que_buffer.put(string)
            else: #sigle process mode
                self._buffer.write(string)

            current_time = time.time()
            if self._start <= 0:
                self._start = current_time
            elif current_time - self._start > self.flush_interval:
                self.flush()

    def writelines(self, sequence):
        if self.pub_socket is None:
            raise ValueError('I/O operation on closed file')
        else:
            for string in sequence:
                self.write(string)

    def _new_buffer(self):
        self._buffer = StringIO()
        self._start = -1