Tools/idle/RemoteInterp.py

import select
import socket
import struct
import sys
import types

VERBOSE = None

class SocketProtocol:
    """A simple protocol for sending strings across a socket"""
    BUF_SIZE = 8192

    def __init__(self, sock):
        self.sock = sock
        self._buffer = ''
        self._closed = 0

    def close(self):
        self._closed = 1
        self.sock.close()

    def send(self, buf):
        """Encode buf and write it on the socket"""
        if VERBOSE:
            VERBOSE.write('send %d:%s\n' % (len(buf), `buf`))
        self.sock.send('%d:%s' % (len(buf), buf))

    def receive(self, timeout=0):
        """Get next complete string from socket or return None

        Raise EOFError on EOF
        """
        buf = self._read_from_buffer()
        if buf is not None:
            return buf
        recvbuf = self._read_from_socket(timeout)
        if recvbuf is None:
            return None
        if recvbuf == '' and self._buffer == '':
            raise EOFError
        if VERBOSE:
            VERBOSE.write('recv %s\n' % `recvbuf`)
        self._buffer = self._buffer + recvbuf
        r = self._read_from_buffer()
        return r

    def _read_from_socket(self, timeout):
        """Does not block"""
        if self._closed:
            return ''
        if timeout is not None:
            r, w, x = select.select([self.sock], [], [], timeout)
        if timeout is None or r:
            return self.sock.recv(self.BUF_SIZE)
        else:
            return None

    def _read_from_buffer(self):
        buf = self._buffer
        i = buf.find(':')
        if i == -1:
            return None
        buflen = int(buf[:i])
        enclen = i + 1 + buflen
        if len(buf) >= enclen:
            s = buf[i+1:enclen]
            self._buffer = buf[enclen:]
            return s
        else:
            self._buffer = buf
        return None

# helpers for registerHandler method below

def get_methods(obj):
    methods = []
    for name in dir(obj):
        attr = getattr(obj, name)
        if callable(attr):
            methods.append(name)
    if type(obj) == types.InstanceType:
        methods = methods + get_methods(obj.__class__)
    if type(obj) == types.ClassType:
        for super in obj.__bases__:
            methods = methods + get_methods(super)
    return methods

class CommandProtocol:
    def __init__(self, sockp):
        self.sockp = sockp
        self.seqno = 0
        self.handlers = {}

    def close(self):
        self.sockp.close()
        self.handlers.clear()

    def registerHandler(self, handler):
        """A Handler is an object with handle_XXX methods"""
        for methname in get_methods(handler):
            if methname[:7] == "handle_":
                name = methname[7:]
                self.handlers[name] = getattr(handler, methname)

    def send(self, cmd, arg='', seqno=None):
        if arg:
            msg = "%s %s" % (cmd, arg)
        else:
            msg = cmd
        if seqno is None:
            seqno = self.get_seqno()
        msgbuf = self.encode_seqno(seqno) + msg
        self.sockp.send(msgbuf)
        if cmd == "reply":
            return
        reply = self.sockp.receive(timeout=None)
        r_cmd, r_arg, r_seqno = self._decode_msg(reply)
        assert r_seqno == seqno and r_cmd == "reply", "bad reply"
        return r_arg

    def _decode_msg(self, msg):
        seqno = self.decode_seqno(msg[:self.SEQNO_ENC_LEN])
        msg = msg[self.SEQNO_ENC_LEN:]
        parts = msg.split(" ", 2)
        if len(parts) == 1:
            cmd = msg
            arg = ''
        else:
            cmd = parts[0]
            arg = parts[1]
        return cmd, arg, seqno

    def dispatch(self):
        msg = self.sockp.receive()
        if msg is None:
            return
        cmd, arg, seqno = self._decode_msg(msg)
        self._current_reply = seqno
        h = self.handlers.get(cmd, self.default_handler)
        try:
            r = h(arg)
        except TypeError, msg:
            raise TypeError, "handle_%s: %s" % (cmd, msg)
        if self._current_reply is None:
            if r is not None:
                sys.stderr.write("ignoring %s return value type %s\n" % \
                                 (cmd, type(r).__name__))
            return
        if r is None:
            r = ''
        if type(r) != types.StringType:
            raise ValueError, "invalid return type for %s" % cmd
        self.send("reply", r, seqno=seqno)

    def reply(self, arg=''):
        """Send a reply immediately

        otherwise reply will be sent when handler returns
        """
        self.send("reply", arg, self._current_reply)
        self._current_reply = None

    def default_handler(self, arg):
        sys.stderr.write("WARNING: unhandled message %s\n" % arg)
        return ''

    SEQNO_ENC_LEN = 4

    def get_seqno(self):
        seqno = self.seqno
        self.seqno = seqno + 1
        return seqno

    def encode_seqno(self, seqno):
        return struct.pack("I", seqno)

    def decode_seqno(self, buf):
        return struct.unpack("I", buf)[0]


class StdioRedirector:
    """Redirect sys.std{in,out,err} to a set of file-like objects"""

    def __init__(self, stdin, stdout, stderr):
        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr

    def redirect(self):
        self.save()
        sys.stdin = self.stdin
        sys.stdout = self.stdout
        sys.stderr = self.stderr

    def save(self):
        self._stdin = sys.stdin
        self._stdout = sys.stdout
        self._stderr = sys.stderr

    def restore(self):
        sys.stdin = self._stdin
        sys.stdout = self._stdout
        sys.stderr = self._stderr

class IOWrapper:
    """Send output from a file-like object across a SocketProtocol

    XXX Should this be more tightly integrated with the CommandProtocol?
    """

    def __init__(self, name, cmdp):
        self.name = name
        self.cmdp = cmdp
        self.buffer = []

class InputWrapper(IOWrapper):
    def write(self, buf):
        # XXX what should this do on Windows?
        raise IOError, (9, '[Errno 9] Bad file descriptor')

    def read(self, arg=None):
        if arg is not None:
            if arg <= 0:
                return ''
        else:
            arg = 0
        return self.cmdp.send(self.name, "read,%s" % arg)

    def readline(self):
        return self.cmdp.send(self.name, "readline")

class OutputWrapper(IOWrapper):
    def write(self, buf):
        self.cmdp.send(self.name, buf)

    def read(self, arg=None):
        return ''

class RemoteInterp:
    def __init__(self, sock):
        self._sock = SocketProtocol(sock)
        self._cmd = CommandProtocol(self._sock)
        self._cmd.registerHandler(self)

    def run(self):
        try:
            while 1:
                self._cmd.dispatch()
        except EOFError:
            pass

    def handle_execfile(self, arg):
        self._cmd.reply()
        io = StdioRedirector(InputWrapper("stdin", self._cmd),
                             OutputWrapper("stdout", self._cmd),
                             OutputWrapper("stderr", self._cmd))
        io.redirect()
        execfile(arg, {'__name__':'__main__'})
        io.restore()
        self._cmd.send("terminated")

    def handle_quit(self, arg):
        self._cmd.reply()
        self._cmd.close()

def startRemoteInterp(id):
    import os
    # UNIX domain sockets are simpler for starters
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.bind("/var/tmp/ri.%s" % id)
    try:
        sock.listen(1)
        cli, addr = sock.accept()
        rinterp = RemoteInterp(cli)
        rinterp.run()
    finally:
        os.unlink("/var/tmp/ri.%s" % id)

class RIClient:
    """Client of the remote interpreter"""
    def __init__(self, sock):
        self._sock = SocketProtocol(sock)
        self._cmd = CommandProtocol(self._sock)
        self._cmd.registerHandler(self)

    def execfile(self, file):
        self._cmd.send("execfile", file)

    def run(self):
        try:
            while 1:
                self._cmd.dispatch()
        except EOFError:
            pass

    def handle_stdout(self, buf):
        sys.stdout.write(buf)
##        sys.stdout.flush()

    def handle_stderr(self, buf):
        sys.stderr.write(buf)

    def handle_stdin(self, arg):
        if arg == "readline":
            return sys.stdin.readline()
        i = arg.find(",") + 1
        bytes = int(arg[i:])
        if bytes == 0:
            return sys.stdin.read()
        else:
            return sys.stdin.read(bytes)

    def handle_terminated(self, arg):
        self._cmd.reply()
        self._cmd.send("quit")
        self._cmd.close()

def riExec(id, file):
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect("/var/tmp/ri.%s" % id)
    cli = RIClient(sock)
    cli.execfile(file)
    cli.run()

if __name__ == "__main__":
    import sys
    import getopt

    SERVER = 1
    opts, args = getopt.getopt(sys.argv[1:], 'cv')
    for o, v in opts:
        if o == '-c':
            SERVER = 0
        elif o == '-v':
            VERBOSE = sys.stderr
    id = args[0]

    if SERVER:
        startRemoteInterp(id)
    else:
        file = args[1]
        riExec(id, file)