Tools/idle/RemoteInterp.py

   1 import select
   2 import socket
   3 import struct
   4 import sys
   5 import types
   6
   7 VERBOSE = None
   8
   9 class SocketProtocol:
  10     """A simple protocol for sending strings across a socket"""
  11     BUF_SIZE = 8192
  12
  13     def __init__(self, sock):
  14         self.sock = sock
  15         self._buffer = ''
  16         self._closed = 0
  17
  18     def close(self):
  19         self._closed = 1
  20         self.sock.close()
  21
  22     def send(self, buf):
  23         """Encode buf and write it on the socket"""
  24         if VERBOSE:
  25             VERBOSE.write('send %d:%s\n' % (len(buf), `buf`))
  26         self.sock.send('%d:%s' % (len(buf), buf))
  27
  28     def receive(self, timeout=0):
  29         """Get next complete string from socket or return None
  30
  31         Raise EOFError on EOF
  32         """
  33         buf = self._read_from_buffer()
  34         if buf is not None:
  35             return buf
  36         recvbuf = self._read_from_socket(timeout)
  37         if recvbuf is None:
  38             return None
  39         if recvbuf == '' and self._buffer == '':
  40             raise EOFError
  41         if VERBOSE:
  42             VERBOSE.write('recv %s\n' % `recvbuf`)
  43         self._buffer = self._buffer + recvbuf
  44         r = self._read_from_buffer()
  45         return r
  46
  47     def _read_from_socket(self, timeout):
  48         """Does not block"""
  49         if self._closed:
  50             return ''
  51         if timeout is not None:
  52             r, w, x = select.select([self.sock], [], [], timeout)
  53         if timeout is None or r:
  54             return self.sock.recv(self.BUF_SIZE)
  55         else:
  56             return None
  57
  58     def _read_from_buffer(self):
  59         buf = self._buffer
  60         i = buf.find(':')
  61         if i == -1:
  62             return None
  63         buflen = int(buf[:i])
  64         enclen = i + 1 + buflen
  65         if len(buf) >= enclen:
  66             s = buf[i+1:enclen]
  67             self._buffer = buf[enclen:]
  68             return s
  69         else:
  70             self._buffer = buf
  71         return None
  72
  73 # helpers for registerHandler method below
  74
  75 def get_methods(obj):
  76     methods = []
  77     for name in dir(obj):
  78         attr = getattr(obj, name)
  79         if callable(attr):
  80             methods.append(name)
  81     if type(obj) == types.InstanceType:
  82         methods = methods + get_methods(obj.__class__)
  83     if type(obj) == types.ClassType:
  84         for super in obj.__bases__:
  85             methods = methods + get_methods(super)
  86     return methods
  87
  88 class CommandProtocol:
  89     def __init__(self, sockp):
  90         self.sockp = sockp
  91         self.seqno = 0
  92         self.handlers = {}
  93
  94     def close(self):
  95         self.sockp.close()
  96         self.handlers.clear()
  97
  98     def registerHandler(self, handler):
  99         """A Handler is an object with handle_XXX methods"""
 100         for methname in get_methods(handler):
 101             if methname[:7] == "handle_":
 102                 name = methname[7:]
 103                 self.handlers[name] = getattr(handler, methname)
 104
 105     def send(self, cmd, arg='', seqno=None):
 106         if arg:
 107             msg = "%s %s" % (cmd, arg)
 108         else:
 109             msg = cmd
 110         if seqno is None:
 111             seqno = self.get_seqno()
 112         msgbuf = self.encode_seqno(seqno) + msg
 113         self.sockp.send(msgbuf)
 114         if cmd == "reply":
 115             return
 116         reply = self.sockp.receive(timeout=None)
 117         r_cmd, r_arg, r_seqno = self._decode_msg(reply)
 118         assert r_seqno == seqno and r_cmd == "reply", "bad reply"
 119         return r_arg
 120
 121     def _decode_msg(self, msg):
 122         seqno = self.decode_seqno(msg[:self.SEQNO_ENC_LEN])
 123         msg = msg[self.SEQNO_ENC_LEN:]
 124         parts = msg.split(" ", 2)
 125         if len(parts) == 1:
 126             cmd = msg
 127             arg = ''
 128         else:
 129             cmd = parts[0]
 130             arg = parts[1]
 131         return cmd, arg, seqno
 132
 133     def dispatch(self):
 134         msg = self.sockp.receive()
 135         if msg is None:
 136             return
 137         cmd, arg, seqno = self._decode_msg(msg)
 138         self._current_reply = seqno
 139         h = self.handlers.get(cmd, self.default_handler)
 140         try:
 141             r = h(arg)
 142         except TypeError, msg:
 143             raise TypeError, "handle_%s: %s" % (cmd, msg)
 144         if self._current_reply is None:
 145             if r is not None:
 146                 sys.stderr.write("ignoring %s return value type %s\n" % \
 147                                  (cmd, type(r).__name__))
 148             return
 149         if r is None:
 150             r = ''
 151         if type(r) != types.StringType:
 152             raise ValueError, "invalid return type for %s" % cmd
 153         self.send("reply", r, seqno=seqno)
 154
 155     def reply(self, arg=''):
 156         """Send a reply immediately
 157
 158         otherwise reply will be sent when handler returns
 159         """
 160         self.send("reply", arg, self._current_reply)
 161         self._current_reply = None
 162
 163     def default_handler(self, arg):
 164         sys.stderr.write("WARNING: unhandled message %s\n" % arg)
 165         return ''
 166
 167     SEQNO_ENC_LEN = 4
 168
 169     def get_seqno(self):
 170         seqno = self.seqno
 171         self.seqno = seqno + 1
 172         return seqno
 173
 174     def encode_seqno(self, seqno):
 175         return struct.pack("I", seqno)
 176
 177     def decode_seqno(self, buf):
 178         return struct.unpack("I", buf)[0]
 179
 180
 181 class StdioRedirector:
 182     """Redirect sys.std{in,out,err} to a set of file-like objects"""
 183
 184     def __init__(self, stdin, stdout, stderr):
 185         self.stdin = stdin
 186         self.stdout = stdout
 187         self.stderr = stderr
 188
 189     def redirect(self):
 190         self.save()
 191         sys.stdin = self.stdin
 192         sys.stdout = self.stdout
 193         sys.stderr = self.stderr
 194
 195     def save(self):
 196         self._stdin = sys.stdin
 197         self._stdout = sys.stdout
 198         self._stderr = sys.stderr
 199
 200     def restore(self):
 201         sys.stdin = self._stdin
 202         sys.stdout = self._stdout
 203         sys.stderr = self._stderr
 204
 205 class IOWrapper:
 206     """Send output from a file-like object across a SocketProtocol
 207
 208     XXX Should this be more tightly integrated with the CommandProtocol?
 209     """
 210
 211     def __init__(self, name, cmdp):
 212         self.name = name
 213         self.cmdp = cmdp
 214         self.buffer = []
 215
 216 class InputWrapper(IOWrapper):
 217     def write(self, buf):
 218         # XXX what should this do on Windows?
 219         raise IOError, (9, '[Errno 9] Bad file descriptor')
 220
 221     def read(self, arg=None):
 222         if arg is not None:
 223             if arg <= 0:
 224                 return ''
 225         else:
 226             arg = 0
 227         return self.cmdp.send(self.name, "read,%s" % arg)
 228
 229     def readline(self):
 230         return self.cmdp.send(self.name, "readline")
 231
 232 class OutputWrapper(IOWrapper):
 233     def write(self, buf):
 234         self.cmdp.send(self.name, buf)
 235
 236     def read(self, arg=None):
 237         return ''
 238
 239 class RemoteInterp:
 240     def __init__(self, sock):
 241         self._sock = SocketProtocol(sock)
 242         self._cmd = CommandProtocol(self._sock)
 243         self._cmd.registerHandler(self)
 244
 245     def run(self):
 246         try:
 247             while 1:
 248                 self._cmd.dispatch()
 249         except EOFError:
 250             pass
 251
 252     def handle_execfile(self, arg):
 253         self._cmd.reply()
 254         io = StdioRedirector(InputWrapper("stdin", self._cmd),
 255                              OutputWrapper("stdout", self._cmd),
 256                              OutputWrapper("stderr", self._cmd))
 257         io.redirect()
 258         execfile(arg, {'__name__':'__main__'})
 259         io.restore()
 260         self._cmd.send("terminated")
 261
 262     def handle_quit(self, arg):
 263         self._cmd.reply()
 264         self._cmd.close()
 265
 266 def startRemoteInterp(id):
 267     import os
 268     # UNIX domain sockets are simpler for starters
 269     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 270     sock.bind("/var/tmp/ri.%s" % id)
 271     try:
 272         sock.listen(1)
 273         cli, addr = sock.accept()
 274         rinterp = RemoteInterp(cli)
 275         rinterp.run()
 276     finally:
 277         os.unlink("/var/tmp/ri.%s" % id)
 278
 279 class RIClient:
 280     """Client of the remote interpreter"""
 281     def __init__(self, sock):
 282         self._sock = SocketProtocol(sock)
 283         self._cmd = CommandProtocol(self._sock)
 284         self._cmd.registerHandler(self)
 285
 286     def execfile(self, file):
 287         self._cmd.send("execfile", file)
 288
 289     def run(self):
 290         try:
 291             while 1:
 292                 self._cmd.dispatch()
 293         except EOFError:
 294             pass
 295
 296     def handle_stdout(self, buf):
 297         sys.stdout.write(buf)
 298 ##        sys.stdout.flush()
 299
 300     def handle_stderr(self, buf):
 301         sys.stderr.write(buf)
 302
 303     def handle_stdin(self, arg):
 304         if arg == "readline":
 305             return sys.stdin.readline()
 306         i = arg.find(",") + 1
 307         bytes = int(arg[i:])
 308         if bytes == 0:
 309             return sys.stdin.read()
 310         else:
 311             return sys.stdin.read(bytes)
 312
 313     def handle_terminated(self, arg):
 314         self._cmd.reply()
 315         self._cmd.send("quit")
 316         self._cmd.close()
 317
 318 def riExec(id, file):
 319     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 320     sock.connect("/var/tmp/ri.%s" % id)
 321     cli = RIClient(sock)
 322     cli.execfile(file)
 323     cli.run()
 324
 325 if __name__ == "__main__":
 326     import sys
 327     import getopt
 328
 329     SERVER = 1
 330     opts, args = getopt.getopt(sys.argv[1:], 'cv')
 331     for o, v in opts:
 332         if o == '-c':
 333             SERVER = 0
 334         elif o == '-v':
 335             VERBOSE = sys.stderr
 336     id = args[0]
 337
 338     if SERVER:
 339         startRemoteInterp(id)
 340     else:
 341         file = args[1]
 342         riExec(id, file)